main
py 413 lines 14.1 KB
Raw
1 from __future__ import annotations
2
3 import dataclasses
4 import importlib.util
5 import inspect
6 import json
7 from concurrent.futures import ThreadPoolExecutor
8 from typing import Any
9
10 import pytest
11
12 if importlib.util.find_spec("scripts.crawl_shard_experiment") is None:
13 pytest.skip("scripts.crawl_shard_experiment is not implemented yet", allow_module_level=True)
14
15 from scripts import crawl_shard_experiment as experiment
16
17
18 def _build_instance(factory: Any, **values: Any) -> Any:
19 signature = inspect.signature(factory)
20 kwargs = {}
21 for name, parameter in signature.parameters.items():
22 if name in values:
23 kwargs[name] = values[name]
24 elif parameter.default is inspect._empty and parameter.kind in (
25 inspect.Parameter.POSITIONAL_OR_KEYWORD,
26 inspect.Parameter.KEYWORD_ONLY,
27 ):
28 raise AssertionError(f"missing required parameter {name!r} for {factory}")
29 return factory(**kwargs)
30
31
32 def _resolve_attr(target: Any, *names: str) -> Any:
33 if isinstance(target, dict):
34 for name in names:
35 if name in target:
36 return target[name]
37 for name in names:
38 if hasattr(target, name):
39 return getattr(target, name)
40 raise AssertionError(f"could not resolve any of {names!r} on {target!r}")
41
42
43 def _resolve_method(target: Any, *names: str):
44 for name in names:
45 candidate = getattr(target, name, None)
46 if callable(candidate):
47 return candidate
48 raise AssertionError(f"could not resolve any of {names!r} on {target!r}")
49
50
51 def _make_tracker(cap: int):
52 return _build_instance(
53 experiment.SharedQuotaTracker,
54 cap=cap,
55 hard_cap=cap,
56 limit=cap,
57 quota_cap=cap,
58 max_calls=cap,
59 )
60
61
62 def _tracker_cap(tracker: Any) -> int:
63 return int(_resolve_attr(tracker, "cap", "hard_cap", "limit", "quota_cap", "max_calls"))
64
65
66 def _tracker_count(tracker: Any) -> int:
67 return int(_resolve_attr(tracker, "count", "used", "api_calls_used", "value", "calls_used"))
68
69
70 def _consume_quota(tracker: Any) -> bool:
71 method = _resolve_method(
72 tracker, "increment", "try_acquire", "acquire", "consume", "record_call"
73 )
74 try:
75 result = method()
76 except Exception as exc: # pragma: no cover
77 message = str(exc).lower()
78 if "quota" in message or "cap" in message or "limit" in message:
79 return False
80 raise
81 if isinstance(result, bool):
82 return result
83 if isinstance(result, int):
84 return result <= _tracker_cap(tracker)
85 return _tracker_count(tracker) <= _tracker_cap(tracker)
86
87
88 def _reset_tracker(tracker: Any) -> None:
89 _resolve_method(tracker, "reset", "clear")()
90
91
92 def _make_budget(seconds: float):
93 return _build_instance(
94 experiment.WallClockBudget,
95 budget_s=seconds,
96 budget_seconds=seconds,
97 seconds=seconds,
98 limit_s=seconds,
99 )
100
101
102 def _budget_exceeded(budget: Any) -> bool:
103 status = _resolve_method(
104 budget, "is_exceeded", "expired", "timed_out", "should_stop", "exhausted"
105 )
106 return bool(status())
107
108
109 def _elapsed_seconds(budget: Any) -> float:
110 for name in ("elapsed_s", "elapsed_seconds", "elapsed"):
111 if hasattr(budget, name):
112 value = getattr(budget, name)
113 return float(value() if callable(value) else value)
114 raise AssertionError(f"could not resolve elapsed time on {budget!r}")
115
116
117 def _make_shard_result(
118 *, repos_found: list[dict[str, Any]], api_calls: int, errors: list[str], wall_clock_s: float
119 ):
120 return _build_instance(
121 experiment.ShardResult,
122 shard_id="github:new-repos:q1",
123 repos_found=repos_found,
124 repos=repos_found,
125 api_calls=api_calls,
126 api_calls_used=api_calls,
127 errors=errors,
128 wall_clock_s=wall_clock_s,
129 duration_s=wall_clock_s,
130 rate_limit_detected=False,
131 rate_limit_events=0,
132 )
133
134
135 def _extract_repos(payload: Any) -> list[dict[str, Any]]:
136 if isinstance(payload, list):
137 return payload
138 for name in ("repos", "repos_found", "items", "results", "merged_repos"):
139 if isinstance(payload, dict) and name in payload:
140 return payload[name]
141 if hasattr(payload, name):
142 return getattr(payload, name)
143 raise AssertionError(f"could not extract merged repos from {payload!r}")
144
145
146 def _metric(result: Any, *names: str) -> Any:
147 return _resolve_attr(result, *names)
148
149
150 def _percent(value: Any) -> float:
151 numeric = float(value)
152 return numeric * 100.0 if abs(numeric) <= 1.0 else numeric
153
154
155 def _canonical_json(payload: Any) -> str:
156 return json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
157
158
159 def _make_report(comparison: dict[str, Any]):
160 report_cls = experiment.ExperimentReport
161 for factory_name in ("from_comparison", "from_metrics", "evaluate", "build"):
162 factory = getattr(report_cls, factory_name, None)
163 if callable(factory):
164 signature = inspect.signature(factory)
165 if "comparison" in signature.parameters:
166 return factory(comparison=comparison)
167 if any(
168 parameter.kind == inspect.Parameter.VAR_KEYWORD
169 for parameter in signature.parameters.values()
170 ):
171 return factory(**comparison)
172 allowed = {
173 name: value for name, value in comparison.items() if name in signature.parameters
174 }
175 return factory(**allowed)
176 signature = inspect.signature(report_cls)
177 if "comparison" in signature.parameters:
178 return report_cls(comparison=comparison)
179 if any(
180 parameter.kind == inspect.Parameter.VAR_KEYWORD
181 for parameter in signature.parameters.values()
182 ):
183 return report_cls(**comparison)
184 allowed = {name: value for name, value in comparison.items() if name in signature.parameters}
185 return report_cls(**allowed)
186
187
188 def _verdict(report: Any) -> str:
189 return str(_resolve_attr(report, "verdict", "status")).lower()
190
191
192 @pytest.fixture
193 def repo_alpha() -> dict[str, Any]:
194 return {
195 "full_name": "octo/alpha",
196 "name": "alpha",
197 "owner": "octo",
198 "stars": 150,
199 "stars_gained": 42,
200 "language": "Python",
201 "topics": ["ai", "testing"],
202 "url": "https://github.com/octo/alpha",
203 }
204
205
206 @pytest.fixture
207 def repo_beta() -> dict[str, Any]:
208 return {
209 "full_name": "octo/beta",
210 "name": "beta",
211 "owner": "octo",
212 "stars": 120,
213 "stars_gained": 15,
214 "language": "Go",
215 "topics": ["infra"],
216 "url": "https://github.com/octo/beta",
217 }
218
219
220 @pytest.fixture
221 def repo_gamma() -> dict[str, Any]:
222 return {
223 "full_name": "tools/gamma",
224 "name": "gamma",
225 "owner": "tools",
226 "stars": 220,
227 "stars_gained": 5,
228 "language": "Rust",
229 "topics": ["developer-tools"],
230 "url": "https://github.com/tools/gamma",
231 }
232
233
234 class TestSharedQuotaTracker:
235 def test_thread_safe_increment_and_cap_enforcement(self) -> None:
236 tracker = _make_tracker(cap=10)
237 with ThreadPoolExecutor(max_workers=8) as pool:
238 successes = list(pool.map(lambda _: _consume_quota(tracker), range(25)))
239 assert sum(bool(result) for result in successes) == 10
240 assert _tracker_count(tracker) == 10
241 assert _consume_quota(tracker) is False
242
243 def test_reset_clears_accumulated_usage(self) -> None:
244 tracker = _make_tracker(cap=3)
245 assert _consume_quota(tracker) is True
246 assert _consume_quota(tracker) is True
247 assert _tracker_count(tracker) == 2
248 _reset_tracker(tracker)
249 assert _tracker_count(tracker) == 0
250 assert _consume_quota(tracker) is True
251
252
253 class TestWallClockBudget:
254 def test_budget_not_exceeded_before_deadline(self, monkeypatch: pytest.MonkeyPatch) -> None:
255 clock = {"now": 100.0}
256 monkeypatch.setattr(experiment.time, "monotonic", lambda: clock["now"])
257 budget = _make_budget(5.0)
258 clock["now"] = 104.9
259 assert _elapsed_seconds(budget) == pytest.approx(4.9)
260 assert _budget_exceeded(budget) is False
261
262 def test_budget_exceeded_triggers_graceful_stop(self, monkeypatch: pytest.MonkeyPatch) -> None:
263 clock = {"now": 250.0}
264 monkeypatch.setattr(experiment.time, "monotonic", lambda: clock["now"])
265 budget = _make_budget(5.0)
266 clock["now"] = 255.1
267 assert _elapsed_seconds(budget) == pytest.approx(5.1)
268 assert _budget_exceeded(budget) is True
269
270
271 def test_shard_result_dataclass_contract(repo_alpha: dict[str, Any]) -> None:
272 assert dataclasses.is_dataclass(experiment.ShardResult)
273 field_names = {field.name for field in dataclasses.fields(experiment.ShardResult)}
274 assert {"repos_found", "api_calls", "errors", "wall_clock_s"}.issubset(field_names)
275 result = _make_shard_result(repos_found=[repo_alpha], api_calls=4, errors=[], wall_clock_s=1.75)
276 assert _resolve_attr(result, "repos_found", "repos") == [repo_alpha]
277 assert _resolve_attr(result, "api_calls", "api_calls_used") == 4
278 assert result.errors == []
279 assert _resolve_attr(result, "wall_clock_s", "duration_s") == pytest.approx(1.75)
280
281
282 class TestDeterministicMerge:
283 def test_deduplicates_by_full_name_and_preserves_star_gain(
284 self, repo_alpha: dict[str, Any], repo_beta: dict[str, Any], repo_gamma: dict[str, Any]
285 ) -> None:
286 first = _make_shard_result(
287 repos_found=[repo_alpha, repo_beta], api_calls=3, errors=[], wall_clock_s=1.1
288 )
289 duplicate_alpha = {**repo_alpha, "stars": 999, "stars_gained": 42}
290 second = _make_shard_result(
291 repos_found=[duplicate_alpha, repo_gamma], api_calls=2, errors=[], wall_clock_s=1.2
292 )
293 merged = experiment.deterministic_merge([first, second])
294 merged_repos = _extract_repos(merged)
295 assert [repo["full_name"] for repo in merged_repos].count("octo/alpha") == 1
296 assert {repo["full_name"] for repo in merged_repos} == {
297 "octo/alpha",
298 "octo/beta",
299 "tools/gamma",
300 }
301 alpha = next(repo for repo in merged_repos if repo["full_name"] == "octo/alpha")
302 assert alpha["stars_gained"] == 42
303
304 def test_same_inputs_produce_byte_identical_output_regardless_of_shard_order(
305 self, repo_alpha: dict[str, Any], repo_beta: dict[str, Any], repo_gamma: dict[str, Any]
306 ) -> None:
307 shard_a = _make_shard_result(
308 repos_found=[repo_alpha, repo_gamma], api_calls=3, errors=[], wall_clock_s=1.0
309 )
310 shard_b = _make_shard_result(
311 repos_found=[repo_beta], api_calls=2, errors=[], wall_clock_s=1.0
312 )
313 merged_ab = experiment.deterministic_merge([shard_a, shard_b])
314 merged_ba = experiment.deterministic_merge([shard_b, shard_a])
315 assert _canonical_json(merged_ab) == _canonical_json(merged_ba)
316
317
318 class TestCompareResults:
319 def test_calculates_speedup_and_api_growth(
320 self, repo_alpha: dict[str, Any], repo_beta: dict[str, Any]
321 ) -> None:
322 canonical = {"repos": [repo_alpha, repo_beta], "metadata": {"note": "stable"}}
323 baseline = {
324 "wall_clock_s": 100.0,
325 "api_calls": 100,
326 "output": canonical,
327 "canonical_output": canonical,
328 }
329 shard = {
330 "wall_clock_s": 70.0,
331 "api_calls": 108,
332 "output": canonical,
333 "canonical_output": canonical,
334 }
335 comparison = experiment.compare_results(baseline, shard)
336 assert _percent(
337 _metric(comparison, "speedup_pct", "speedup_percent", "speedup")
338 ) == pytest.approx(30.0)
339 assert _percent(
340 _metric(comparison, "api_growth_pct", "api_growth_percent", "api_growth")
341 ) == pytest.approx(8.0)
342
343 def test_output_stability_ignores_timestamp_fields(
344 self, repo_alpha: dict[str, Any], repo_beta: dict[str, Any]
345 ) -> None:
346 baseline = {
347 "wall_clock_s": 100.0,
348 "api_calls": 100,
349 "output": {
350 "week": "2026-W24",
351 "crawled_at": "2026-06-13T12:00:00Z",
352 "report_generated_at": "2026-06-13T12:00:01Z",
353 "repos": [repo_alpha, repo_beta],
354 },
355 }
356 shard = {
357 "wall_clock_s": 75.0,
358 "api_calls": 105,
359 "output": {
360 "week": "2026-W24",
361 "crawled_at": "2026-06-13T12:10:00Z",
362 "report_generated_at": "2026-06-13T12:10:01Z",
363 "repos": [repo_alpha, repo_beta],
364 },
365 }
366 comparison = experiment.compare_results(baseline, shard)
367 assert (
368 bool(_metric(comparison, "output_stable", "stable_output", "is_output_stable")) is True
369 )
370
371
372 class TestExperimentReport:
373 def test_pass_verdict_when_all_guardrails_are_met(self) -> None:
374 report = _make_report(
375 {
376 "speedup_pct": 30.0,
377 "api_growth_pct": 8.0,
378 "rate_limit_regression": False,
379 "output_stable": True,
380 "partial_data": False,
381 "baseline_complete": True,
382 "shard_complete": True,
383 }
384 )
385 assert _verdict(report) == "pass"
386
387 def test_fail_verdict_when_any_required_criterion_fails(self) -> None:
388 report = _make_report(
389 {
390 "speedup_pct": 20.0,
391 "api_growth_pct": 8.0,
392 "rate_limit_regression": False,
393 "output_stable": True,
394 "partial_data": False,
395 "baseline_complete": True,
396 "shard_complete": True,
397 }
398 )
399 assert _verdict(report) == "fail"
400
401 def test_inconclusive_verdict_for_partial_data(self) -> None:
402 report = _make_report(
403 {
404 "speedup_pct": None,
405 "api_growth_pct": None,
406 "rate_limit_regression": False,
407 "output_stable": False,
408 "partial_data": True,
409 "baseline_complete": True,
410 "shard_complete": False,
411 }
412 )
413 assert _verdict(report) == "inconclusive"