gabriel / musehub public
test_quota_via_refs.py python
257 lines 9.9 KB
Raw
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65 refactor: enforce gRPC framing on all MWP wire traffic Sonnet 4.6 minor ⚠ breaking 156 days ago
1 """TDD: quota queries use musehub_object_refs, not musehub_objects.repo_id.
2
3 Phase 4 of the object-refs architecture: storage accounting must walk the refs
4 table so that content-addressed dedup is correctly attributed. Two repos that
5 share an object must each pay for it independently (the object exists once in
6 storage but is intentionally charged to each repo that references it).
7
8 Coverage matrix:
9 1. wire_push quota: shared object is counted for both repos independently.
10 2. wire_push quota: dedup push to same repo is not double-counted.
11 3. wire_push quota: exceeded quota is rejected.
12 4. get_repo_stats: total_objects and total_size_bytes use refs join.
13 5. get_repo_stats: shared object counted once per repo.
14 """
15 from __future__ import annotations
16
17 import hashlib
18 import uuid
19 import zlib
20 from unittest.mock import AsyncMock, MagicMock, patch
21
22 import msgpack
23 import pytest
24 import sqlalchemy as sa
25 from httpx import AsyncClient
26 from sqlalchemy.ext.asyncio import AsyncSession
27
28 from musehub.db import musehub_models as db
29 from musehub.types.json_types import JSONValue, StrDict
30 from tests.factories import create_repo
31
32
33 def _oid(raw: bytes) -> str:
34 return "sha256:" + hashlib.sha256(raw).hexdigest()
35
36
37 def _mp(data: JSONValue) -> bytes:
38 return msgpack.packb(data, use_bin_type=True)
39
40
41 def _zlib(raw: bytes) -> bytes:
42 return zlib.compress(raw, level=1)
43
44
45 # ---------------------------------------------------------------------------
46 # Test 1: shared object is counted for both repos independently
47 # ---------------------------------------------------------------------------
48
49 @pytest.mark.asyncio
50 async def test_quota_counts_shared_object_per_repo(
51 client: AsyncClient,
52 db_session: AsyncSession,
53 wire_headers: StrDict,
54 ) -> None:
55 """Two repos pushing the same bytes each count that object toward their quota.
56
57 This is correct accounting: both repos reference the object and would each
58 need to pay for it if the other repo were deleted.
59 """
60 raw = b"# shared object for quota test\n" * 50
61 oid = _oid(raw)
62
63 repo_a = await create_repo(db_session, slug=f"quota-shared-a-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
64 repo_b = await create_repo(db_session, slug=f"quota-shared-b-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
65
66 # Push to repo_a first
67 r_a = await client.post(
68 f"/{repo_a.owner}/{repo_a.slug}/push/object-pack",
69 content=_mp({"objects": [{"object_id": oid, "content": _zlib(raw), "encoding": "zlib", "path": "shared.md"}]}),
70 headers=wire_headers,
71 )
72 assert r_a.status_code == 200
73
74 # Push same bytes to repo_b
75 r_b = await client.post(
76 f"/{repo_b.owner}/{repo_b.slug}/push/object-pack",
77 content=_mp({"objects": [{"object_id": oid, "content": _zlib(raw), "encoding": "zlib", "path": "shared.md"}]}),
78 headers=wire_headers,
79 )
80 assert r_b.status_code == 200
81
82 # Both repos should have a ref
83 for repo in (repo_a, repo_b):
84 ref = (await db_session.execute(
85 sa.select(db.MusehubObjectRef).where(
86 db.MusehubObjectRef.repo_id == repo.repo_id,
87 db.MusehubObjectRef.object_id == oid,
88 )
89 )).scalar_one_or_none()
90 assert ref is not None, f"repo {repo.slug} must have a ref row"
91
92 # Only one row in musehub_objects (dedup invariant)
93 obj_count = (await db_session.execute(
94 sa.select(sa.func.count()).select_from(db.MusehubObject).where(
95 db.MusehubObject.object_id == oid
96 )
97 )).scalar_one()
98 assert obj_count == 1, "dedup: one object row"
99
100 # Quota query via refs: each repo counts the object
101 def _quota_for(repo_id: str):
102 return (
103 sa.select(sa.func.coalesce(sa.func.sum(db.MusehubObject.size_bytes), 0))
104 .join(db.MusehubObjectRef, db.MusehubObject.object_id == db.MusehubObjectRef.object_id)
105 .where(db.MusehubObjectRef.repo_id == repo_id)
106 )
107
108 size_a = (await db_session.execute(_quota_for(repo_a.repo_id))).scalar_one()
109 size_b = (await db_session.execute(_quota_for(repo_b.repo_id))).scalar_one()
110
111 assert size_a > 0, "repo_a quota must count the shared object"
112 assert size_b > 0, "repo_b quota must count the shared object"
113 assert size_a == size_b, "both repos pay the same amount for the shared object"
114
115
116 # ---------------------------------------------------------------------------
117 # Test 2: re-pushing same object to same repo does not double-count quota
118 # ---------------------------------------------------------------------------
119
120 @pytest.mark.asyncio
121 async def test_quota_idempotent_push_no_double_count(
122 client: AsyncClient,
123 db_session: AsyncSession,
124 wire_headers: StrDict,
125 ) -> None:
126 """Pushing the same object twice to the same repo must not inflate quota."""
127 raw = b"# idempotent quota test object\n" * 40
128 oid = _oid(raw)
129 repo = await create_repo(db_session, slug=f"quota-idem-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
130
131 url = f"/{repo.owner}/{repo.slug}/push/object-pack"
132 payload = _mp({"objects": [{"object_id": oid, "content": _zlib(raw), "encoding": "zlib", "path": "idem.md"}]})
133
134 for _ in range(3):
135 r = await client.post(url, content=payload, headers=wire_headers)
136 assert r.status_code == 200
137
138 ref_count = (await db_session.execute(
139 sa.select(sa.func.count()).select_from(db.MusehubObjectRef).where(
140 db.MusehubObjectRef.repo_id == repo.repo_id,
141 db.MusehubObjectRef.object_id == oid,
142 )
143 )).scalar_one()
144 assert ref_count == 1, "idempotent push must not create duplicate refs"
145
146 quota_used = (await db_session.execute(
147 sa.select(sa.func.coalesce(sa.func.sum(db.MusehubObject.size_bytes), 0))
148 .join(db.MusehubObjectRef, db.MusehubObject.object_id == db.MusehubObjectRef.object_id)
149 .where(db.MusehubObjectRef.repo_id == repo.repo_id)
150 )).scalar_one()
151
152 # Must equal exactly one copy of the object's size
153 obj_size = (await db_session.execute(
154 sa.select(db.MusehubObject.size_bytes).where(db.MusehubObject.object_id == oid)
155 )).scalar_one()
156
157 assert quota_used == obj_size, \
158 f"quota must count the object once; got {quota_used} vs obj_size={obj_size}"
159
160
161 # ---------------------------------------------------------------------------
162 # Test 3: push rejected when quota exceeded
163 # ---------------------------------------------------------------------------
164
165 @pytest.mark.asyncio
166 async def test_push_rejected_when_quota_exceeded(
167 db_session: AsyncSession,
168 ) -> None:
169 """wire_push must reject pushes that would exceed per_repo_quota_bytes."""
170 from musehub.services.musehub_wire import wire_push
171 from musehub.models.wire import WirePushRequest, WireBundle, WireObject
172
173 raw = b"x" * 100
174 oid = _oid(raw)
175 repo = await create_repo(db_session, slug=f"quota-reject-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
176
177 bundle = WireBundle(
178 objects=[WireObject(object_id=oid, path="big.md", content=raw)],
179 commits=[],
180 snapshots=[],
181 )
182 req = WirePushRequest(repo_id=repo.repo_id, branch="main", bundle=bundle)
183
184 with patch("musehub.services.musehub_wire.settings") as mock_settings:
185 mock_settings.per_repo_quota_bytes = 1 # 1 byte — always exceeded
186 result = await wire_push(db_session, repo.repo_id, req, pusher_id="test-user-wire")
187
188 assert result.ok is False
189 assert "quota" in result.message.lower(), \
190 f"rejection message must mention quota; got: {result.message}"
191
192
193 # ---------------------------------------------------------------------------
194 # Test 4: get_repo_stats counts and sizes via refs join
195 # ---------------------------------------------------------------------------
196
197 @pytest.mark.asyncio
198 async def test_repo_stats_uses_refs_join(
199 db_session: AsyncSession,
200 ) -> None:
201 """get_repo_stats total_objects and total_size_bytes must count via refs."""
202 from musehub.services.musehub_repository import get_repo_home_stats
203
204 repo = await create_repo(db_session, slug=f"stats-refs-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
205
206 # Insert an object + ref directly
207 oid = _oid(b"stats test object content")
208 obj = db.MusehubObject(
209 object_id=oid,
210 path="stats.md",
211 size_bytes=42,
212 disk_path="",
213 )
214 db_session.add(obj)
215 db_session.add(db.MusehubObjectRef(repo_id=repo.repo_id, object_id=oid))
216 await db_session.commit()
217
218 stats = await get_repo_home_stats(db_session, repo.repo_id, ref="main")
219
220 assert stats["total_objects"] == 1
221 assert stats["total_size_bytes"] == 42
222
223
224 # ---------------------------------------------------------------------------
225 # Test 5: get_repo_stats shared object counted once per repo
226 # ---------------------------------------------------------------------------
227
228 @pytest.mark.asyncio
229 async def test_repo_stats_shared_object_per_repo(
230 db_session: AsyncSession,
231 ) -> None:
232 """A shared object must appear in each repo's stats independently."""
233 from musehub.services.musehub_repository import get_repo_home_stats
234
235 repo_a = await create_repo(db_session, slug=f"stats-shared-a-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
236 repo_b = await create_repo(db_session, slug=f"stats-shared-b-{uuid.uuid4().hex[:8]}", owner="test-user-wire")
237
238 oid = _oid(b"shared stats object")
239 obj = db.MusehubObject(
240 object_id=oid,
241 path="shared.md",
242 size_bytes=100,
243 disk_path="",
244 )
245 db_session.add(obj)
246 # Both repos hold a ref
247 db_session.add(db.MusehubObjectRef(repo_id=repo_a.repo_id, object_id=oid))
248 db_session.add(db.MusehubObjectRef(repo_id=repo_b.repo_id, object_id=oid))
249 await db_session.commit()
250
251 stats_a = await get_repo_home_stats(db_session, repo_a.repo_id, ref="main")
252 stats_b = await get_repo_home_stats(db_session, repo_b.repo_id, ref="main")
253
254 assert stats_a["total_objects"] == 1
255 assert stats_a["total_size_bytes"] == 100
256 assert stats_b["total_objects"] == 1
257 assert stats_b["total_size_bytes"] == 100
File History 1 commit
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65 refactor: enforce gRPC framing on all MWP wire traffic Sonnet 4.6 minor ⚠ 156 days ago