| 1 | """Tests for scripts/context_budget.py.""" |
| 2 | |
| 3 | import sys |
| 4 | from datetime import datetime, timedelta, timezone |
| 5 | from pathlib import Path |
| 6 | |
| 7 | sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
| 8 | |
| 9 | from scripts.context_budget import ( |
| 10 | assemble_historical_context, |
| 11 | compress_stale_trends, |
| 12 | compress_to_budget, |
| 13 | keep_top_noise_patterns, |
| 14 | prune_stale_predictions, |
| 15 | word_count, |
| 16 | ) |
| 17 | |
| 18 | # --- word_count tests --- |
| 19 | |
| 20 | |
| 21 | class TestWordCount: |
| 22 | def test_empty_string(self): |
| 23 | assert word_count("") == 0 |
| 24 | |
| 25 | def test_single_word(self): |
| 26 | assert word_count("hello") == 1 |
| 27 | |
| 28 | def test_multiple_words(self): |
| 29 | assert word_count("one two three four five") == 5 |
| 30 | |
| 31 | def test_multiline(self): |
| 32 | text = "line one\nline two\nline three" |
| 33 | assert word_count(text) == 6 |
| 34 | |
| 35 | def test_extra_whitespace(self): |
| 36 | assert word_count(" hello world ") == 2 |
| 37 | |
| 38 | |
| 39 | # --- compress_to_budget tests --- |
| 40 | |
| 41 | |
| 42 | class TestCompressToBudget: |
| 43 | def test_empty_text(self): |
| 44 | assert compress_to_budget("", 100) == "" |
| 45 | |
| 46 | def test_whitespace_only(self): |
| 47 | assert compress_to_budget(" \n ", 100) == "" |
| 48 | |
| 49 | def test_within_budget(self): |
| 50 | text = "This is a short sentence." |
| 51 | result = compress_to_budget(text, 100) |
| 52 | assert result == text.strip() |
| 53 | assert word_count(result) <= 100 |
| 54 | |
| 55 | def test_exceeds_budget_truncates(self): |
| 56 | text = " ".join(f"word{i}" for i in range(100)) |
| 57 | result = compress_to_budget(text, 10) |
| 58 | assert word_count(result) <= 11 # 10 + possible trailing "..." |
| 59 | assert result.endswith("...") |
| 60 | |
| 61 | def test_zero_budget(self): |
| 62 | assert compress_to_budget("hello world", 0) == "" |
| 63 | |
| 64 | def test_negative_budget(self): |
| 65 | assert compress_to_budget("hello world", -5) == "" |
| 66 | |
| 67 | def test_preserves_line_structure(self): |
| 68 | text = "line one two\nline three four five\nline six seven eight nine ten" |
| 69 | result = compress_to_budget(text, 5) |
| 70 | # Should keep first line (3 words) and partial of second |
| 71 | assert word_count(result) <= 6 # budget + trailing word |
| 72 | assert "..." in result or word_count(result) <= 5 |
| 73 | |
| 74 | def test_budget_enforcement_strict(self): |
| 75 | text = " ".join(["word"] * 1000) |
| 76 | result = compress_to_budget(text, 50) |
| 77 | # Word count should not greatly exceed budget |
| 78 | assert word_count(result) <= 51 |
| 79 | |
| 80 | |
| 81 | # --- prune_stale_predictions tests --- |
| 82 | |
| 83 | |
| 84 | class TestPruneStale: |
| 85 | def test_drops_old_unconfirmed(self): |
| 86 | now = datetime(2026, 6, 12, tzinfo=timezone.utc) |
| 87 | text = "- [2026-03-01] prediction about AI growth\n- [2026-06-01] recent prediction" |
| 88 | result = prune_stale_predictions(text, now=now) |
| 89 | assert "2026-03-01" not in result |
| 90 | assert "2026-06-01" in result |
| 91 | |
| 92 | def test_keeps_confirmed(self): |
| 93 | now = datetime(2026, 6, 12, tzinfo=timezone.utc) |
| 94 | text = "- [2026-01-01] old but confirmed ✓" |
| 95 | result = prune_stale_predictions(text, now=now) |
| 96 | assert "2026-01-01" in result |
| 97 | |
| 98 | def test_keeps_recent(self): |
| 99 | now = datetime(2026, 6, 12, tzinfo=timezone.utc) |
| 100 | recent = (now - timedelta(weeks=2)).strftime("%Y-%m-%d") |
| 101 | text = f"- [{recent}] recent prediction" |
| 102 | result = prune_stale_predictions(text, now=now) |
| 103 | assert recent in result |
| 104 | |
| 105 | def test_empty_text(self): |
| 106 | assert prune_stale_predictions("") == "" |
| 107 | |
| 108 | |
| 109 | # --- compress_stale_trends tests --- |
| 110 | |
| 111 | |
| 112 | class TestCompressStaleTrends: |
| 113 | def test_compresses_stale_trend(self): |
| 114 | now = datetime(2026, 6, 12, tzinfo=timezone.utc) |
| 115 | text = ( |
| 116 | "## Trend: Blockchain hype\n" |
| 117 | "- [2026-01-01] some old signal\n" |
| 118 | "- [2026-01-15] another old signal\n" |
| 119 | ) |
| 120 | result = compress_stale_trends(text, now=now) |
| 121 | assert "Blockchain hype: no recent signal" in result |
| 122 | assert "some old signal" not in result |
| 123 | |
| 124 | def test_keeps_active_trend(self): |
| 125 | now = datetime(2026, 6, 12, tzinfo=timezone.utc) |
| 126 | recent = (now - timedelta(days=5)).strftime("%Y-%m-%d") |
| 127 | text = f"## Trend: AI growth\n- [{recent}] strong signal this week\n" |
| 128 | result = compress_stale_trends(text, now=now) |
| 129 | assert "strong signal this week" in result |
| 130 | |
| 131 | def test_empty_text(self): |
| 132 | assert compress_stale_trends("") == "" |
| 133 | |
| 134 | |
| 135 | # --- keep_top_noise_patterns tests --- |
| 136 | |
| 137 | |
| 138 | class TestNoisePatterns: |
| 139 | def test_keeps_only_top_n(self): |
| 140 | text = ( |
| 141 | "## Noise\n" |
| 142 | "- noise pattern 1\n" |
| 143 | "- noise pattern 2\n" |
| 144 | "- noise pattern 3\n" |
| 145 | "- noise pattern 4\n" |
| 146 | "- noise pattern 5\n" |
| 147 | ) |
| 148 | result = keep_top_noise_patterns(text, max_patterns=3) |
| 149 | assert "noise pattern 1" in result |
| 150 | assert "noise pattern 2" in result |
| 151 | assert "noise pattern 3" in result |
| 152 | assert "noise pattern 4" not in result |
| 153 | assert "noise pattern 5" not in result |
| 154 | |
| 155 | def test_fewer_than_max(self): |
| 156 | text = "## Noise\n- pattern 1\n- pattern 2\n" |
| 157 | result = keep_top_noise_patterns(text, max_patterns=3) |
| 158 | assert "pattern 1" in result |
| 159 | assert "pattern 2" in result |
| 160 | |
| 161 | def test_non_noise_sections_untouched(self): |
| 162 | text = "## Signals\n- sig 1\n- sig 2\n- sig 3\n- sig 4\n- sig 5\n" |
| 163 | result = keep_top_noise_patterns(text, max_patterns=3) |
| 164 | assert "sig 5" in result # Not filtered |
| 165 | |
| 166 | |
| 167 | # --- assemble_historical_context tests --- |
| 168 | |
| 169 | |
| 170 | class TestAssemble: |
| 171 | def test_assembles_from_files(self, tmp_path): |
| 172 | rolling = tmp_path / "rolling.md" |
| 173 | yearly = tmp_path / "yearly.md" |
| 174 | prev_week = tmp_path / "prev_week.md" |
| 175 | |
| 176 | rolling.write_text("Rolling context with some words here today.") |
| 177 | yearly.write_text("Yearly summary of important events and trends.") |
| 178 | prev_week.write_text("Last week's highlights and observations.") |
| 179 | |
| 180 | result = assemble_historical_context( |
| 181 | rolling_path=str(rolling), |
| 182 | yearly_path=str(yearly), |
| 183 | prev_week_path=str(prev_week), |
| 184 | max_total_words=1500, |
| 185 | ) |
| 186 | |
| 187 | assert "## Rolling Context" in result |
| 188 | assert "## Yearly Context" in result |
| 189 | assert "## Previous Week" in result |
| 190 | |
| 191 | def test_handles_missing_files(self, tmp_path): |
| 192 | result = assemble_historical_context( |
| 193 | rolling_path=str(tmp_path / "nonexistent.md"), |
| 194 | yearly_path=None, |
| 195 | prev_week_path=None, |
| 196 | max_total_words=1500, |
| 197 | ) |
| 198 | # Should not crash, returns empty or partial |
| 199 | assert isinstance(result, str) |
| 200 | |
| 201 | def test_all_none_paths(self): |
| 202 | result = assemble_historical_context( |
| 203 | rolling_path=None, |
| 204 | yearly_path=None, |
| 205 | prev_week_path=None, |
| 206 | max_total_words=1500, |
| 207 | ) |
| 208 | assert result == "" |
| 209 | |
| 210 | def test_budget_enforcement(self, tmp_path): |
| 211 | # Create a large file that exceeds budget |
| 212 | large_text = " ".join(["word"] * 2000) |
| 213 | rolling = tmp_path / "rolling.md" |
| 214 | rolling.write_text(large_text) |
| 215 | |
| 216 | yearly = tmp_path / "yearly.md" |
| 217 | yearly.write_text(large_text) |
| 218 | |
| 219 | result = assemble_historical_context( |
| 220 | rolling_path=str(rolling), |
| 221 | yearly_path=str(yearly), |
| 222 | max_total_words=100, |
| 223 | ) |
| 224 | # Total words should respect the budget (with small margin for headers) |
| 225 | assert word_count(result) <= 110 |
| 226 | |
| 227 | def test_output_is_markdown(self, tmp_path): |
| 228 | rolling = tmp_path / "rolling.md" |
| 229 | rolling.write_text("Some rolling context data points here.") |
| 230 | |
| 231 | result = assemble_historical_context( |
| 232 | rolling_path=str(rolling), |
| 233 | max_total_words=1500, |
| 234 | ) |
| 235 | # Should contain markdown headers |
| 236 | assert result.startswith("## ") |
| 237 | |
| 238 | def test_budget_allocation_scaling(self, tmp_path): |
| 239 | """Verify budgets scale proportionally with max_total_words.""" |
| 240 | text = " ".join(["word"] * 1000) |
| 241 | rolling = tmp_path / "rolling.md" |
| 242 | rolling.write_text(text) |
| 243 | |
| 244 | # With max_total_words=750 (half of default 1500), |
| 245 | # rolling budget should be ~250 (half of 500) |
| 246 | result = assemble_historical_context( |
| 247 | rolling_path=str(rolling), |
| 248 | max_total_words=750, |
| 249 | ) |
| 250 | # Section content should be roughly half the default rolling budget |
| 251 | content = result.replace("## Rolling Context\n\n", "") |
| 252 | assert word_count(content) <= 280 # ~250 + margin |
| 253 | |
| 254 | def test_pruning_applied_during_assembly(self, tmp_path): |
| 255 | now = datetime(2026, 6, 12, tzinfo=timezone.utc) |
| 256 | rolling = tmp_path / "rolling.md" |
| 257 | rolling.write_text( |
| 258 | "- [2026-01-01] stale prediction from January\n" |
| 259 | "- [2026-06-10] fresh prediction from this week\n" |
| 260 | ) |
| 261 | |
| 262 | result = assemble_historical_context( |
| 263 | rolling_path=str(rolling), |
| 264 | max_total_words=1500, |
| 265 | now=now, |
| 266 | ) |
| 267 | assert "2026-01-01" not in result |
| 268 | assert "2026-06-10" in result |