| 1 | import unittest |
| 2 | |
| 3 | from lib import fusion, schema |
| 4 | |
| 5 | |
| 6 | def make_item(item_id: str, source: str, url: str, title: str, rank_score: float) -> schema.SourceItem: |
| 7 | return schema.SourceItem( |
| 8 | item_id=item_id, |
| 9 | source=source, |
| 10 | title=title, |
| 11 | body=title, |
| 12 | url=url, |
| 13 | relevance_hint=rank_score, |
| 14 | snippet=title, |
| 15 | metadata={ |
| 16 | "local_relevance": rank_score, |
| 17 | "freshness": 80, |
| 18 | "engagement_score": 5, |
| 19 | "source_quality": 0.7, |
| 20 | }, |
| 21 | ) |
| 22 | |
| 23 | |
| 24 | class FusionV3Tests(unittest.TestCase): |
| 25 | def test_weighted_rrf_merges_duplicate_urls(self): |
| 26 | plan = schema.QueryPlan( |
| 27 | intent="breaking_news", |
| 28 | freshness_mode="strict_recent", |
| 29 | cluster_mode="story", |
| 30 | raw_topic="test", |
| 31 | subqueries=[ |
| 32 | schema.SubQuery(label="primary", search_query="test", ranking_query="What happened in test?", sources=["reddit", "x"], weight=0.7), |
| 33 | schema.SubQuery(label="reaction", search_query="test reaction", ranking_query="What are the reactions to test?", sources=["x"], weight=0.3), |
| 34 | ], |
| 35 | source_weights={"reddit": 0.4, "x": 0.6}, |
| 36 | ) |
| 37 | shared = "https://example.com/shared" |
| 38 | streams = { |
| 39 | ("primary", "reddit"): [make_item("r1", "reddit", shared, "Shared item", 0.8)], |
| 40 | ("primary", "x"): [make_item("x1", "x", shared, "Shared item", 0.9)], |
| 41 | ("reaction", "x"): [make_item("x2", "x", "https://example.com/unique", "Unique item", 0.7)], |
| 42 | } |
| 43 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=10) |
| 44 | self.assertEqual(2, len(candidates)) |
| 45 | merged = next(candidate for candidate in candidates if candidate.url == shared) |
| 46 | self.assertEqual({"primary"}, set(merged.subquery_labels)) |
| 47 | self.assertEqual(2, len(merged.native_ranks)) |
| 48 | self.assertEqual({"reddit", "x"}, set(merged.sources)) |
| 49 | self.assertEqual(2, len(merged.source_items)) |
| 50 | |
| 51 | def test_diversify_pool_guarantees_min_per_qualifying_source(self): |
| 52 | """Every qualifying source (local_relevance >= 0.25) gets at least 2 |
| 53 | items in the fused pool. |
| 54 | |
| 55 | Dominant sources (x, tiktok) get high weights, so pure-RRF truncation |
| 56 | would squeeze out low-weight sources entirely. The diversity guarantee |
| 57 | must reserve at least 2 slots per qualifying active source. All sources |
| 58 | here have rank_score=0.8 (well above the 0.25 threshold), so every |
| 59 | source qualifies for reserved slots. |
| 60 | """ |
| 61 | sources = ["reddit", "hackernews", "x", "tiktok", "bluesky", "youtube"] |
| 62 | # Heavily skewed weights: x and tiktok dominate. |
| 63 | weights = { |
| 64 | "x": 3.0, |
| 65 | "tiktok": 2.5, |
| 66 | "reddit": 0.5, |
| 67 | "hackernews": 0.4, |
| 68 | "bluesky": 0.3, |
| 69 | "youtube": 0.3, |
| 70 | } |
| 71 | plan = schema.QueryPlan( |
| 72 | intent="concept", |
| 73 | freshness_mode="relaxed", |
| 74 | cluster_mode="concept", |
| 75 | raw_topic="RAG", |
| 76 | subqueries=[ |
| 77 | schema.SubQuery( |
| 78 | label="primary", |
| 79 | search_query="RAG", |
| 80 | ranking_query="What is RAG?", |
| 81 | sources=sources, |
| 82 | weight=1.0, |
| 83 | ), |
| 84 | ], |
| 85 | source_weights=weights, |
| 86 | ) |
| 87 | streams: dict[tuple[str, str], list[schema.SourceItem]] = {} |
| 88 | for src in sources: |
| 89 | items = [] |
| 90 | for rank in range(4): |
| 91 | items.append( |
| 92 | make_item( |
| 93 | item_id=f"{src}_{rank}", |
| 94 | source=src, |
| 95 | url=f"https://{src}.example.com/{rank}", |
| 96 | title=f"{src} item {rank}", |
| 97 | rank_score=0.8, |
| 98 | ) |
| 99 | ) |
| 100 | streams[("primary", src)] = items |
| 101 | |
| 102 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=12) |
| 103 | self.assertEqual(12, len(candidates)) |
| 104 | |
| 105 | source_counts: dict[str, int] = {} |
| 106 | for c in candidates: |
| 107 | source_counts[c.source] = source_counts.get(c.source, 0) + 1 |
| 108 | |
| 109 | for src in sources: |
| 110 | self.assertGreaterEqual( |
| 111 | source_counts.get(src, 0), |
| 112 | 2, |
| 113 | f"Source '{src}' has {source_counts.get(src, 0)} items, expected >= 2", |
| 114 | ) |
| 115 | |
| 116 | def test_diversify_pool_denies_slots_for_low_relevance_source(self): |
| 117 | """Sources with best local_relevance < 0.25 do not get reserved slots. |
| 118 | |
| 119 | Create two sources: 'x' with local_relevance=0.5 (qualifies) and |
| 120 | 'reddit' with local_relevance=0.1 (below threshold). With a tight |
| 121 | pool_limit, the high-relevance source gets reserved slots while |
| 122 | the low-relevance source must compete on RRF merit alone. |
| 123 | """ |
| 124 | plan = schema.QueryPlan( |
| 125 | intent="concept", |
| 126 | freshness_mode="relaxed", |
| 127 | cluster_mode="concept", |
| 128 | raw_topic="test", |
| 129 | subqueries=[ |
| 130 | schema.SubQuery( |
| 131 | label="primary", |
| 132 | search_query="test", |
| 133 | ranking_query="What is test?", |
| 134 | sources=["x", "reddit"], |
| 135 | weight=1.0, |
| 136 | ), |
| 137 | ], |
| 138 | source_weights={"x": 1.0, "reddit": 1.0}, |
| 139 | ) |
| 140 | |
| 141 | # x items: high relevance (0.5) -- qualifies for diversity reservation |
| 142 | x_items = [ |
| 143 | make_item(f"x_{i}", "x", f"https://x.example.com/{i}", f"x item {i}", 0.5) |
| 144 | for i in range(4) |
| 145 | ] |
| 146 | |
| 147 | # reddit items: low relevance (0.1) -- below threshold, no reserved slots |
| 148 | reddit_items = [ |
| 149 | make_item(f"r_{i}", "reddit", f"https://reddit.example.com/{i}", f"reddit item {i}", 0.1) |
| 150 | for i in range(4) |
| 151 | ] |
| 152 | |
| 153 | streams = { |
| 154 | ("primary", "x"): x_items, |
| 155 | ("primary", "reddit"): reddit_items, |
| 156 | } |
| 157 | |
| 158 | # pool_limit=3: x gets 2 reserved + 1 more by RRF. Reddit has no |
| 159 | # reserved slots, so it must out-score x items in the remainder. |
| 160 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=3) |
| 161 | self.assertEqual(3, len(candidates)) |
| 162 | |
| 163 | # x must have at least 2 (reserved slots) |
| 164 | x_count = sum(1 for c in candidates if c.source == "x") |
| 165 | self.assertGreaterEqual(x_count, 2, "x should have at least 2 reserved slots") |
| 166 | |
| 167 | def test_diversify_pool_no_reservation_when_all_below_threshold(self): |
| 168 | """When all sources are below the relevance threshold, no reserved slots |
| 169 | are granted. The pool is filled purely by RRF score order.""" |
| 170 | plan = schema.QueryPlan( |
| 171 | intent="concept", |
| 172 | freshness_mode="relaxed", |
| 173 | cluster_mode="concept", |
| 174 | raw_topic="test", |
| 175 | subqueries=[ |
| 176 | schema.SubQuery( |
| 177 | label="primary", |
| 178 | search_query="test", |
| 179 | ranking_query="What is test?", |
| 180 | sources=["x", "reddit", "hackernews"], |
| 181 | weight=1.0, |
| 182 | ), |
| 183 | ], |
| 184 | # Give x a much higher weight so its items get higher RRF scores |
| 185 | source_weights={"x": 3.0, "reddit": 0.3, "hackernews": 0.3}, |
| 186 | ) |
| 187 | |
| 188 | streams: dict[tuple[str, str], list[schema.SourceItem]] = {} |
| 189 | # All sources below threshold (local_relevance = 0.1) |
| 190 | for src in ["x", "reddit", "hackernews"]: |
| 191 | items = [ |
| 192 | make_item(f"{src}_{i}", src, f"https://{src}.example.com/{i}", f"{src} item {i}", 0.1) |
| 193 | for i in range(4) |
| 194 | ] |
| 195 | streams[("primary", src)] = items |
| 196 | |
| 197 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=4) |
| 198 | self.assertEqual(4, len(candidates)) |
| 199 | |
| 200 | # With no diversity reservation and x having 3x the weight, |
| 201 | # x should dominate the top slots purely on RRF score |
| 202 | source_counts: dict[str, int] = {} |
| 203 | for c in candidates: |
| 204 | source_counts[c.source] = source_counts.get(c.source, 0) + 1 |
| 205 | |
| 206 | # x has 3x weight so its RRF scores are ~3x higher than reddit/hn. |
| 207 | # All 4 x items should beat all reddit/hackernews items. |
| 208 | self.assertEqual( |
| 209 | source_counts.get("x", 0), |
| 210 | 4, |
| 211 | f"Expected x to take all 4 slots on pure RRF merit, got {source_counts}", |
| 212 | ) |
| 213 | |
| 214 | def test_diversify_pool_threshold_boundary(self): |
| 215 | """Source with best local_relevance exactly at the threshold (0.25) |
| 216 | qualifies for reserved slots.""" |
| 217 | plan = schema.QueryPlan( |
| 218 | intent="concept", |
| 219 | freshness_mode="relaxed", |
| 220 | cluster_mode="concept", |
| 221 | raw_topic="boundary", |
| 222 | subqueries=[ |
| 223 | schema.SubQuery( |
| 224 | label="primary", |
| 225 | search_query="boundary", |
| 226 | ranking_query="What is boundary?", |
| 227 | sources=["x", "reddit"], |
| 228 | weight=1.0, |
| 229 | ), |
| 230 | ], |
| 231 | # Give x much higher weight so it would dominate without reservation |
| 232 | source_weights={"x": 5.0, "reddit": 0.1}, |
| 233 | ) |
| 234 | |
| 235 | x_items = [ |
| 236 | make_item(f"x_{i}", "x", f"https://x.example.com/{i}", f"x item {i}", 0.8) |
| 237 | for i in range(6) |
| 238 | ] |
| 239 | |
| 240 | # reddit at exactly the threshold |
| 241 | reddit_items = [ |
| 242 | make_item(f"r_{i}", "reddit", f"https://reddit.example.com/{i}", f"reddit item {i}", 0.25) |
| 243 | for i in range(3) |
| 244 | ] |
| 245 | |
| 246 | streams = { |
| 247 | ("primary", "x"): x_items, |
| 248 | ("primary", "reddit"): reddit_items, |
| 249 | } |
| 250 | |
| 251 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=6) |
| 252 | self.assertEqual(6, len(candidates)) |
| 253 | |
| 254 | reddit_count = sum(1 for c in candidates if c.source == "reddit") |
| 255 | self.assertGreaterEqual( |
| 256 | reddit_count, |
| 257 | 2, |
| 258 | f"reddit (local_relevance=0.25, at threshold) should get 2 reserved slots, got {reddit_count}", |
| 259 | ) |
| 260 | |
| 261 | |
| 262 | def make_item_with_author( |
| 263 | item_id: str, source: str, url: str, title: str, rank_score: float, author: str | None = None, |
| 264 | ) -> schema.SourceItem: |
| 265 | return schema.SourceItem( |
| 266 | item_id=item_id, |
| 267 | source=source, |
| 268 | title=title, |
| 269 | body=title, |
| 270 | url=url, |
| 271 | author=author, |
| 272 | relevance_hint=rank_score, |
| 273 | snippet=title, |
| 274 | metadata={ |
| 275 | "local_relevance": rank_score, |
| 276 | "freshness": 80, |
| 277 | "engagement_score": 5, |
| 278 | "source_quality": 0.7, |
| 279 | }, |
| 280 | ) |
| 281 | |
| 282 | |
| 283 | class TestPerAuthorCap(unittest.TestCase): |
| 284 | """Per-author cap: no single author should have more than 3 items in fused pool.""" |
| 285 | |
| 286 | def _make_plan(self, sources: list[str]) -> schema.QueryPlan: |
| 287 | return schema.QueryPlan( |
| 288 | intent="breaking_news", |
| 289 | freshness_mode="strict_recent", |
| 290 | cluster_mode="story", |
| 291 | raw_topic="test", |
| 292 | subqueries=[ |
| 293 | schema.SubQuery( |
| 294 | label="primary", |
| 295 | search_query="test", |
| 296 | ranking_query="test", |
| 297 | sources=sources, |
| 298 | weight=1.0, |
| 299 | ), |
| 300 | ], |
| 301 | source_weights={s: 1.0 for s in sources}, |
| 302 | ) |
| 303 | |
| 304 | def test_author_with_8_items_capped_to_3(self): |
| 305 | """@grok scenario: 8 items from the same author, only best 3 survive.""" |
| 306 | plan = self._make_plan(["x"]) |
| 307 | items = [ |
| 308 | make_item_with_author( |
| 309 | f"x_{i}", "x", f"https://x.com/{i}", f"grok summary {i}", 0.7, author="@grok", |
| 310 | ) |
| 311 | for i in range(8) |
| 312 | ] |
| 313 | streams = {("primary", "x"): items} |
| 314 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 315 | grok_count = sum( |
| 316 | 1 for c in candidates |
| 317 | if any(si.author == "@grok" for si in c.source_items) |
| 318 | ) |
| 319 | self.assertLessEqual(grok_count, 3, f"@grok should be capped at 3, got {grok_count}") |
| 320 | |
| 321 | def test_author_with_3_items_all_kept(self): |
| 322 | """Author with exactly 3 items should keep all of them.""" |
| 323 | plan = self._make_plan(["x"]) |
| 324 | items = [ |
| 325 | make_item_with_author( |
| 326 | f"x_{i}", "x", f"https://x.com/{i}", f"author3 post {i}", 0.7, author="@author3", |
| 327 | ) |
| 328 | for i in range(3) |
| 329 | ] |
| 330 | streams = {("primary", "x"): items} |
| 331 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 332 | count = sum( |
| 333 | 1 for c in candidates |
| 334 | if any(si.author == "@author3" for si in c.source_items) |
| 335 | ) |
| 336 | self.assertEqual(count, 3) |
| 337 | |
| 338 | def test_items_without_author_not_capped(self): |
| 339 | """Items with no author field should never be dropped by the cap.""" |
| 340 | plan = self._make_plan(["reddit"]) |
| 341 | items = [ |
| 342 | make_item_with_author( |
| 343 | f"r_{i}", "reddit", f"https://reddit.com/{i}", f"post {i}", 0.7, author=None, |
| 344 | ) |
| 345 | for i in range(6) |
| 346 | ] |
| 347 | streams = {("primary", "reddit"): items} |
| 348 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 349 | self.assertEqual(len(candidates), 6) |
| 350 | |
| 351 | def test_multiple_authors_capped_independently(self): |
| 352 | """Two prolific authors each get capped to 3 independently.""" |
| 353 | plan = self._make_plan(["x"]) |
| 354 | items = [] |
| 355 | for i in range(5): |
| 356 | items.append(make_item_with_author( |
| 357 | f"grok_{i}", "x", f"https://x.com/grok/{i}", f"grok {i}", 0.7, author="@grok", |
| 358 | )) |
| 359 | for i in range(5): |
| 360 | items.append(make_item_with_author( |
| 361 | f"spam_{i}", "x", f"https://x.com/spam/{i}", f"spam {i}", 0.6, author="@spammer", |
| 362 | )) |
| 363 | streams = {("primary", "x"): items} |
| 364 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 365 | grok_count = sum(1 for c in candidates if any(si.author == "@grok" for si in c.source_items)) |
| 366 | spam_count = sum(1 for c in candidates if any(si.author == "@spammer" for si in c.source_items)) |
| 367 | self.assertLessEqual(grok_count, 3) |
| 368 | self.assertLessEqual(spam_count, 3) |
| 369 | |
| 370 | def test_cap_keeps_best_items_by_rrf_order(self): |
| 371 | """The cap should keep the first (highest-ranked) items per author.""" |
| 372 | plan = self._make_plan(["x"]) |
| 373 | # Items with decreasing relevance scores so ranking is deterministic |
| 374 | items = [ |
| 375 | make_item_with_author( |
| 376 | f"x_{i}", "x", f"https://x.com/{i}", f"post {i}", 0.9 - (i * 0.05), author="@prolific", |
| 377 | ) |
| 378 | for i in range(5) |
| 379 | ] |
| 380 | streams = {("primary", "x"): items} |
| 381 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 382 | kept_ids = {c.item_id for c in candidates if any(si.author == "@prolific" for si in c.source_items)} |
| 383 | # The top 3 items (x_0, x_1, x_2) should be kept |
| 384 | self.assertLessEqual(len(kept_ids), 3) |
| 385 | |
| 386 | |
| 387 | class TestUrlNormalization(unittest.TestCase): |
| 388 | def test_strips_www(self): |
| 389 | from lib.fusion import _normalize_url |
| 390 | self.assertEqual( |
| 391 | _normalize_url("https://www.reddit.com/r/test"), |
| 392 | _normalize_url("https://reddit.com/r/test"), |
| 393 | ) |
| 394 | |
| 395 | def test_strips_old_prefix(self): |
| 396 | from lib.fusion import _normalize_url |
| 397 | self.assertEqual( |
| 398 | _normalize_url("https://old.reddit.com/r/test"), |
| 399 | _normalize_url("https://reddit.com/r/test"), |
| 400 | ) |
| 401 | |
| 402 | def test_strips_mobile_prefix(self): |
| 403 | from lib.fusion import _normalize_url |
| 404 | self.assertEqual( |
| 405 | _normalize_url("https://m.youtube.com/watch?v=abc"), |
| 406 | _normalize_url("https://youtube.com/watch?v=abc"), |
| 407 | ) |
| 408 | |
| 409 | def test_strips_utm_params(self): |
| 410 | from lib.fusion import _normalize_url |
| 411 | self.assertEqual( |
| 412 | _normalize_url("https://example.com/page?utm_source=twitter&id=5"), |
| 413 | _normalize_url("https://example.com/page?id=5"), |
| 414 | ) |
| 415 | |
| 416 | def test_strips_trailing_slash(self): |
| 417 | from lib.fusion import _normalize_url |
| 418 | self.assertEqual( |
| 419 | _normalize_url("https://example.com/page/"), |
| 420 | _normalize_url("https://example.com/page"), |
| 421 | ) |
| 422 | |
| 423 | def test_preserves_non_tracking_params(self): |
| 424 | from lib.fusion import _normalize_url |
| 425 | result = _normalize_url("https://example.com/page?id=5&sort=new") |
| 426 | self.assertIn("id=5", result) |
| 427 | self.assertIn("sort=new", result) |
| 428 | |
| 429 | def test_case_insensitive(self): |
| 430 | from lib.fusion import _normalize_url |
| 431 | self.assertEqual( |
| 432 | _normalize_url("https://Reddit.com/r/Test"), |
| 433 | _normalize_url("https://reddit.com/r/test"), |
| 434 | ) |
| 435 | |
| 436 | if __name__ == "__main__": |
| 437 | unittest.main() |
| 438 | |
| 439 | |
| 440 | class OutOfWindowSortTests(unittest.TestCase): |
| 441 | def _candidate(self, name: str, published_at: str | None, confidence: str, rrf: float) -> schema.Candidate: |
| 442 | item = schema.SourceItem( |
| 443 | item_id=name, |
| 444 | source="youtube", |
| 445 | title=name, |
| 446 | body="body", |
| 447 | url=f"https://youtube.com/watch?v={name}", |
| 448 | published_at=published_at, |
| 449 | date_confidence=confidence, |
| 450 | ) |
| 451 | return schema.Candidate( |
| 452 | candidate_id=name, |
| 453 | item_id=name, |
| 454 | source="youtube", |
| 455 | title=name, |
| 456 | url=item.url, |
| 457 | snippet="snippet", |
| 458 | subquery_labels=["primary"], |
| 459 | native_ranks={"primary:youtube": 1}, |
| 460 | local_relevance=0.9, |
| 461 | freshness=90, |
| 462 | engagement=60.0, |
| 463 | source_quality=0.85, |
| 464 | rrf_score=rrf, |
| 465 | source_items=[item], |
| 466 | ) |
| 467 | |
| 468 | def test_out_of_window_sorts_below_in_window(self): |
| 469 | stale = self._candidate("stale", "2025-10-15", "low", rrf=0.9) |
| 470 | fresh = self._candidate("fresh", "2026-07-20", "high", rrf=0.01) |
| 471 | ordered = sorted([stale, fresh], key=fusion._candidate_sort_key) |
| 472 | self.assertEqual(["fresh", "stale"], [c.candidate_id for c in ordered]) |
| 473 | |
| 474 | def test_undated_candidate_keeps_its_place(self): |
| 475 | undated = self._candidate("undated", None, "low", rrf=0.9) |
| 476 | dated = self._candidate("dated", "2026-07-20", "high", rrf=0.01) |
| 477 | ordered = sorted([dated, undated], key=fusion._candidate_sort_key) |
| 478 | self.assertEqual(["undated", "dated"], [c.candidate_id for c in ordered]) |
| 479 | |
| 480 | def test_adapter_high_confidence_outside_window_is_still_stale(self): |
| 481 | """Adapters may supply date_confidence='high' for old dates. |
| 482 | |
| 483 | The window membership must be derived from the actual date compared to |
| 484 | the run window, not solely from adapter-provided date_confidence. |
| 485 | """ |
| 486 | item = schema.SourceItem( |
| 487 | item_id="old_job", |
| 488 | source="jobs", |
| 489 | title="Old job posting", |
| 490 | body="body", |
| 491 | url="https://example.com/job", |
| 492 | published_at="2025-10-15", |
| 493 | date_confidence="high", |
| 494 | ) |
| 495 | stale_with_high_confidence = schema.Candidate( |
| 496 | candidate_id="stale_high", |
| 497 | item_id="old_job", |
| 498 | source="jobs", |
| 499 | title="Old job posting", |
| 500 | url="https://example.com/job", |
| 501 | snippet="snippet", |
| 502 | subquery_labels=["primary"], |
| 503 | native_ranks={"primary:jobs": 1}, |
| 504 | local_relevance=0.9, |
| 505 | freshness=90, |
| 506 | engagement=60.0, |
| 507 | source_quality=0.85, |
| 508 | rrf_score=0.9, |
| 509 | source_items=[item], |
| 510 | metadata={"range_from": "2026-06-15", "range_to": "2026-07-15"}, |
| 511 | ) |
| 512 | fresh = self._candidate("fresh", "2026-07-10", "high", rrf=0.01) |
| 513 | fresh.metadata["range_from"] = "2026-06-15" |
| 514 | fresh.metadata["range_to"] = "2026-07-15" |
| 515 | |
| 516 | self.assertTrue(schema.candidate_out_of_window(stale_with_high_confidence)) |
| 517 | self.assertFalse(schema.candidate_out_of_window(fresh)) |
| 518 | |
| 519 | ordered = sorted([stale_with_high_confidence, fresh], key=fusion._candidate_sort_key) |
| 520 | self.assertEqual(["fresh", "stale_high"], [c.candidate_id for c in ordered]) |
| 521 | |
| 522 | |
| 523 | def _comments(n: int) -> list[dict]: |
| 524 | return [{"score": 100 - i, "author": f"u{i}", "excerpt": f"comment {i}"} for i in range(n)] |
| 525 | |
| 526 | |
| 527 | class FusionEnrichedCopyTests(unittest.TestCase): |
| 528 | """Same thread, two subquery streams, per-stream ids collide (R1 / R1): |
| 529 | the candidate must keep the copy that carries the enrichment.""" |
| 530 | |
| 531 | def _plan(self): |
| 532 | return schema.QueryPlan( |
| 533 | intent="breaking_news", |
| 534 | freshness_mode="strict_recent", |
| 535 | cluster_mode="story", |
| 536 | raw_topic="kanye west", |
| 537 | subqueries=[ |
| 538 | schema.SubQuery(label="primary", search_query="kanye west", ranking_query="kanye west", sources=["reddit"], weight=1.0), |
| 539 | schema.SubQuery(label="russia", search_query="kanye russia", ranking_query="kanye russia", sources=["reddit"], weight=0.8), |
| 540 | ], |
| 541 | source_weights={}, |
| 542 | ) |
| 543 | |
| 544 | def test_enriched_second_copy_is_kept(self): |
| 545 | url = "https://www.reddit.com/r/Music/comments/1vy0ilk/kanye_wests_soldout_russia/" |
| 546 | bare = make_item("R1", "reddit", url, "Kanye West's Russia shows canceled", 0.6) |
| 547 | rich = make_item("R1", "reddit", url, "Kanye West's Russia shows canceled", 0.3) |
| 548 | rich.metadata["top_comments"] = _comments(10) |
| 549 | rich.metadata["comment_insights"] = ["Putin doesn't care"] |
| 550 | rich.engagement = {"score": 20859, "num_comments": 741} |
| 551 | bare.engagement = {"score": 18941, "num_comments": 713} |
| 552 | streams = {("primary", "reddit"): [bare], ("russia", "reddit"): [rich]} |
| 553 | |
| 554 | candidates = fusion.weighted_rrf(streams, self._plan(), pool_limit=10) |
| 555 | |
| 556 | self.assertEqual(1, len(candidates)) |
| 557 | cand = candidates[0] |
| 558 | self.assertEqual(1, len(cand.source_items)) |
| 559 | self.assertEqual(10, len(cand.source_items[0].metadata["top_comments"])) |
| 560 | self.assertEqual(["Putin doesn't care"], cand.source_items[0].metadata["comment_insights"]) |
| 561 | self.assertEqual(20859, cand.source_items[0].engagement["score"]) |
| 562 | self.assertEqual({"primary:reddit", "russia:reddit"}, set(cand.native_ranks)) |
| 563 | |
| 564 | def test_enriched_first_copy_is_retained(self): |
| 565 | url = "https://www.reddit.com/r/Music/comments/1vy0ilk/kanye_wests_soldout_russia/" |
| 566 | rich = make_item("R1", "reddit", url, "Kanye West's Russia shows canceled", 0.6) |
| 567 | rich.metadata["top_comments"] = _comments(4) |
| 568 | bare = make_item("R1", "reddit", url, "Kanye West's Russia shows canceled", 0.3) |
| 569 | streams = {("primary", "reddit"): [rich], ("russia", "reddit"): [bare]} |
| 570 | |
| 571 | candidates = fusion.weighted_rrf(streams, self._plan(), pool_limit=10) |
| 572 | |
| 573 | self.assertEqual(4, len(candidates[0].source_items[0].metadata["top_comments"])) |
| 574 | |
| 575 | def test_distinct_urls_with_colliding_ids_stay_separate(self): |
| 576 | a = make_item("R1", "reddit", "https://www.reddit.com/r/Kanye/comments/aaa/one/", "Thread one", 0.6) |
| 577 | b = make_item("R1", "reddit", "https://www.reddit.com/r/Kanye/comments/bbb/two/", "Thread two", 0.5) |
| 578 | streams = {("primary", "reddit"): [a], ("russia", "reddit"): [b]} |
| 579 | |
| 580 | candidates = fusion.weighted_rrf(streams, self._plan(), pool_limit=10) |
| 581 | |
| 582 | self.assertEqual(2, len(candidates)) |
| 583 | |
| 584 | def test_collapse_duplicate_urls_merges_enrichment(self): |
| 585 | url = "https://www.reddit.com/r/Music/comments/1vy0ilk/kanye/" |
| 586 | bare = make_item("R1", "reddit", url, "Kanye", 0.6) |
| 587 | rich = make_item("R1", "reddit", url, "Kanye", 0.3) |
| 588 | rich.metadata["top_comments"] = _comments(3) |
| 589 | other = make_item("R2", "reddit", "https://www.reddit.com/r/Kanye/comments/zzz/other/", "Other", 0.4) |
| 590 | |
| 591 | out = fusion.collapse_duplicate_urls([bare, rich, other]) |
| 592 | |
| 593 | self.assertEqual(2, len(out)) |
| 594 | self.assertIs(out[0], bare) |
| 595 | self.assertEqual(3, len(out[0].metadata["top_comments"])) |
| 596 | self.assertIs(out[1], other) |
| 597 |