main
py 363 lines 12.3 KB
Raw
1 """Tests for momentum_tracker and calibrate_hype_risk."""
2
3 import json
4 import sys
5 from pathlib import Path
6
7 import pytest
8
9 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts"))
10 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts" / "archive"))
11 from momentum_tracker import ( # noqa: E402
12 classify_momentum,
13 compute_decay_rate,
14 current_iso_week,
15 extract_correlated_repos,
16 find_correlation_file,
17 get_repo_stars_gained,
18 iso_week_to_date,
19 load_json_safe,
20 run_momentum_tracking,
21 track_repo_momentum,
22 update_predictions_validated,
23 week_offset,
24 )
25 from calibrate_hype_risk import ( # noqa: E402
26 build_actual_outcomes,
27 build_predictions,
28 compute_calibration,
29 generate_recommendations,
30 risk_to_expected_outcome,
31 run_calibration,
32 )
33
34
35 # ---------------------------------------------------------------------------
36 # momentum_tracker tests
37 # ---------------------------------------------------------------------------
38
39
40 class TestWeekUtils:
41 """Test ISO week utility functions."""
42
43 def test_current_iso_week_format(self):
44 """current_iso_week returns YYYY-WNN format."""
45 week = current_iso_week()
46 assert len(week) >= 7
47 assert "-W" in week
48
49 def test_week_offset_forward(self):
50 """week_offset moves forward correctly."""
51 assert week_offset("2026-W21", 2) == "2026-W23"
52
53 def test_week_offset_backward(self):
54 """week_offset moves backward correctly."""
55 assert week_offset("2026-W03", -2) == "2026-W01"
56
57 def test_week_offset_year_boundary(self):
58 """week_offset crosses year boundary."""
59 result = week_offset("2025-W52", 2)
60 assert result.startswith("2026-W")
61
62 def test_iso_week_to_date(self):
63 """iso_week_to_date returns correct Monday."""
64 dt = iso_week_to_date("2026-W01")
65 assert dt.weekday() == 0 # Monday
66
67
68 class TestClassifyMomentum:
69 """Test momentum classification logic."""
70
71 def test_sustained_strong_growth(self):
72 """Still gaining 20%+ of initial → sustained."""
73 assert classify_momentum(100, 50, 30, lag=4) == "sustained"
74
75 def test_faded_no_growth(self):
76 """Zero growth at checkpoints → faded."""
77 assert classify_momentum(100, 5, 3, lag=4) == "faded"
78
79 def test_faded_zero_initial(self):
80 """Zero initial gained → faded."""
81 assert classify_momentum(0, 10, 5, lag=4) == "faded"
82
83 def test_sustained_week2_only(self):
84 """With lag=2, only week2 data used."""
85 assert classify_momentum(100, 30, None, lag=2) == "sustained"
86
87 def test_faded_no_data(self):
88 """No follow-up data → faded (conservative)."""
89 assert classify_momentum(100, None, None, lag=4) == "faded"
90
91 def test_sustained_at_threshold(self):
92 """Exactly 20% of initial → sustained."""
93 assert classify_momentum(100, 20, 20, lag=4) == "sustained"
94
95 def test_faded_below_threshold(self):
96 """Just below 20% → faded."""
97 assert classify_momentum(100, 19, 19, lag=4) == "faded"
98
99
100 class TestDecayRate:
101 """Test decay rate computation."""
102
103 def test_no_decay(self):
104 assert compute_decay_rate(100, 100) == 0.0
105
106 def test_full_decay(self):
107 assert compute_decay_rate(100, 0) == 1.0
108
109 def test_partial_decay(self):
110 assert compute_decay_rate(100, 50) == 0.5
111
112 def test_zero_initial(self):
113 assert compute_decay_rate(0, 50) == 0.0
114
115 def test_negative_clamped(self):
116 assert compute_decay_rate(100, 150) == 0.0
117
118
119 class TestExtractCorrelatedRepos:
120 """Test extraction from correlation data."""
121
122 def test_extracts_correlated(self):
123 data = {
124 "correlations": [
125 {"repo": "org/a", "press_correlated": True},
126 {"repo": "org/b", "press_correlated": False},
127 {"repo": "org/c", "press_correlated": True},
128 ]
129 }
130 result = extract_correlated_repos(data)
131 assert len(result) == 2
132 assert result[0]["repo"] == "org/a"
133 assert result[1]["repo"] == "org/c"
134
135 def test_empty_correlations(self):
136 assert extract_correlated_repos({"correlations": []}) == []
137
138
139 class TestGetRepoStarsGained:
140 """Test star extraction from raw data."""
141
142 def test_from_repos_key(self):
143 data = {"repos": [{"full_name": "org/x", "stars_gained": 42}]}
144 assert get_repo_stars_gained(data, "org/x") == 42
145
146 def test_missing_repo(self):
147 data = {"repos": [{"full_name": "org/x", "stars_gained": 42}]}
148 assert get_repo_stars_gained(data, "org/y") is None
149
150 def test_none_data(self):
151 assert get_repo_stars_gained(None, "org/x") is None
152
153
154 class TestTrackRepoMomentum:
155 """Test single repo momentum tracking."""
156
157 def test_with_data(self, tmp_path):
158 w2_data = {"repos": [{"full_name": "org/x", "stars_gained": 50}]}
159 (tmp_path / "2026-W23.json").write_text(json.dumps(w2_data))
160 w4_data = {"repos": [{"full_name": "org/x", "stars_gained": 30}]}
161 (tmp_path / "2026-W25.json").write_text(json.dumps(w4_data))
162
163 result = track_repo_momentum("org/x", 200, tmp_path, "2026-W21", lag=4)
164 assert result["repo"] == "org/x"
165 assert result["initial_stars_gained"] == 200
166 assert result["week2_stars_gained"] == 50
167 assert result["week4_stars_gained"] == 30
168 assert result["classification"] == "faded"
169 assert result["decay_rate"] > 0
170
171 def test_missing_weeks(self, tmp_path):
172 result = track_repo_momentum("org/x", 100, tmp_path, "2026-W21", lag=4)
173 assert result["week2_stars_gained"] is None
174 assert result["week4_stars_gained"] is None
175 assert result["classification"] == "faded"
176
177
178 class TestUpdatePredictions:
179 """Test predictions.jsonl update."""
180
181 def test_updates_matching(self, tmp_path):
182 preds = [
183 {"repo": "org/a", "prediction": "rising_star", "week": "2026-W20"},
184 {"repo": "org/b", "prediction": "rising_star", "week": "2026-W20"},
185 ]
186 pred_path = tmp_path / "predictions.jsonl"
187 pred_path.write_text("\n".join(json.dumps(p) for p in preds) + "\n")
188
189 tracked = [
190 {"repo": "org/a", "classification": "sustained"},
191 {"repo": "org/b", "classification": "faded"},
192 ]
193 updated = update_predictions_validated(pred_path, tracked)
194 assert updated == 2
195
196 lines = pred_path.read_text().strip().split("\n")
197 result = [json.loads(line) for line in lines]
198 assert result[0]["validated"] is True
199 assert result[1]["validated"] is False
200
201 def test_skips_already_validated(self, tmp_path):
202 preds = [{"repo": "org/a", "validated": True, "week": "2026-W20"}]
203 pred_path = tmp_path / "predictions.jsonl"
204 pred_path.write_text(json.dumps(preds[0]) + "\n")
205
206 tracked = [{"repo": "org/a", "classification": "faded"}]
207 updated = update_predictions_validated(pred_path, tracked)
208 assert updated == 0
209
210 def test_missing_file(self, tmp_path):
211 updated = update_predictions_validated(tmp_path / "nope.jsonl", [])
212 assert updated == 0
213
214
215 class TestRunMomentumTracking:
216 """Integration test for full tracking run."""
217
218 def test_no_correlations(self, tmp_path, monkeypatch):
219 import scripts.topic_paths as tp
220
221 monkeypatch.setattr(tp, "DATA_ROOT", tmp_path / "data")
222 (tmp_path / "data" / "raw").mkdir(parents=True)
223 (tmp_path / "data" / "analyzed").mkdir(parents=True)
224 (tmp_path / "data" / "metrics").mkdir(parents=True)
225
226 result = run_momentum_tracking(topic_id=None, week="2026-W21", lag=4)
227 assert result["week"] == "2026-W21"
228 assert result["tracked_repos"] == []
229 assert result["summary"]["total"] == 0
230
231
232 # ---------------------------------------------------------------------------
233 # calibrate_hype_risk tests
234 # ---------------------------------------------------------------------------
235
236
237 class TestBuildActualOutcomes:
238 """Test outcome extraction from momentum data."""
239
240 def test_extracts_outcomes(self):
241 data = [
242 {
243 "tracked_repos": [
244 {"repo": "org/a", "classification": "sustained"},
245 {"repo": "org/b", "classification": "faded"},
246 ]
247 }
248 ]
249 outcomes = build_actual_outcomes(data)
250 assert outcomes["org/a"] == "sustained"
251 assert outcomes["org/b"] == "faded"
252
253 def test_latest_wins(self):
254 data = [
255 {"tracked_repos": [{"repo": "org/a", "classification": "faded"}]},
256 {"tracked_repos": [{"repo": "org/a", "classification": "sustained"}]},
257 ]
258 outcomes = build_actual_outcomes(data)
259 assert outcomes["org/a"] == "sustained"
260
261
262 class TestBuildPredictions:
263 """Test prediction extraction from hype risk data."""
264
265 def test_extracts_risks(self):
266 data = [
267 {
268 "assessments": [
269 {"repo": "org/a", "hype_risk": "high"},
270 {"repo": "org/b", "hype_risk": "low"},
271 ]
272 }
273 ]
274 preds = build_predictions(data)
275 assert preds["org/a"] == "high"
276 assert preds["org/b"] == "low"
277
278
279 class TestRiskToExpectedOutcome:
280 """Test risk level to outcome mapping."""
281
282 def test_high_expects_faded(self):
283 assert risk_to_expected_outcome("high") == "faded"
284
285 def test_low_expects_sustained(self):
286 assert risk_to_expected_outcome("low") == "sustained"
287
288 def test_very_low_expects_sustained(self):
289 assert risk_to_expected_outcome("very_low") == "sustained"
290
291 def test_medium_uncertain(self):
292 assert risk_to_expected_outcome("medium") is None
293
294 def test_none_uncertain(self):
295 assert risk_to_expected_outcome("none") is None
296
297
298 class TestComputeCalibration:
299 """Test calibration computation."""
300
301 def test_perfect_accuracy(self):
302 predictions = {"org/a": "high", "org/b": "low"}
303 actuals = {"org/a": "faded", "org/b": "sustained"}
304 result = compute_calibration(predictions, actuals)
305 assert result["samples"] == 2
306 assert result["accuracy_by_category"]["high"]["accuracy"] == 1.0
307 assert result["accuracy_by_category"]["low"]["accuracy"] == 1.0
308
309 def test_partial_accuracy(self):
310 predictions = {"org/a": "high", "org/b": "high"}
311 actuals = {"org/a": "faded", "org/b": "sustained"}
312 result = compute_calibration(predictions, actuals)
313 assert result["samples"] == 2
314 assert result["accuracy_by_category"]["high"]["correct"] == 1
315 assert result["accuracy_by_category"]["high"]["predicted"] == 2
316
317 def test_no_overlap(self):
318 predictions = {"org/a": "high"}
319 actuals = {"org/z": "faded"}
320 result = compute_calibration(predictions, actuals)
321 assert result["samples"] == 0
322
323
324 class TestGenerateRecommendations:
325 """Test recommendation generation."""
326
327 def test_low_high_accuracy_triggers_adjustment(self):
328 calibration = {
329 "accuracy_by_category": {
330 "high": {"predicted": 10, "correct": 5, "accuracy": 0.5}
331 }
332 }
333 actuals = {"org/a": "sustained", "org/b": "faded"}
334 recs = generate_recommendations(calibration, actuals)
335 params = [r["parameter"] for r in recs]
336 assert "high_risk_decay_threshold" in params
337
338 def test_good_accuracy_no_changes(self):
339 calibration = {
340 "accuracy_by_category": {
341 "high": {"predicted": 10, "correct": 9, "accuracy": 0.9},
342 "low": {"predicted": 10, "correct": 9, "accuracy": 0.9},
343 }
344 }
345 actuals = {"org/a": "sustained", "org/b": "faded"}
346 recs = generate_recommendations(calibration, actuals)
347 params = [r["parameter"] for r in recs]
348 assert "no_changes" in params
349
350
351 class TestRunCalibration:
352 """Integration test for calibration."""
353
354 def test_no_data(self, tmp_path, monkeypatch):
355 import scripts.topic_paths as tp
356
357 monkeypatch.setattr(tp, "DATA_ROOT", tmp_path / "data")
358 (tmp_path / "data" / "metrics").mkdir(parents=True)
359 (tmp_path / "data" / "analyzed").mkdir(parents=True)
360
361 result = run_calibration(topic_id=None, output_path=str(tmp_path / "out.json"))
362 assert result["samples"] == 0
363 assert result["recommended_adjustments"] == []