"""Unit tests for scripts/entity_resolution_v2.py (BET 1b: blocking + embedding similarity + LLM adjudication — dry-run tooling). Pure-python: no DB, no network, no embeddings endpoint, no LLM. The HTTP embedding backend and the Anthropic adjudicator are NEVER called here — tests use fakes. Importable without psycopg (both scripts guard the driver import). Run: pytest packages/memory-engine-v2/tests/test_entity_resolution_v2.py """ from __future__ import annotations import importlib.util import json import subprocess import sys from pathlib import Path _SCRIPTS = Path(__file__).resolve().parent.parent / "scripts" def _load(name: str): if name in sys.modules: return sys.modules[name] spec = importlib.util.spec_from_file_location(name, _SCRIPTS / f"{name}.py") assert spec and spec.loader mod = importlib.util.module_from_spec(spec) sys.modules[name] = mod # register before exec (py3.13+ dataclasses) spec.loader.exec_module(mod) return mod er2 = _load("entity_resolution_v2") v1 = _load("backfill_entity_reconciliation") def make_entity(eid: str, name: str, aliases: list[str] | None = None, facts: int = 0, rels: int = 0) -> "v1.Entity": return v1.Entity( id=eid, canonical_name=name, aliases=aliases or [], provenance_event_ids=[], fact_count=facts, rel_count=rels, norm_forms={v1._normalize_surface(name), *(v1._normalize_surface(a) for a in (aliases or []))}, ) class FakeAdjudicator(er2.Adjudicator): """Deterministic verdicts keyed by frozenset of entity ids.""" def __init__(self, verdicts: dict[frozenset, str]) -> None: self.verdicts = verdicts self.calls: list[frozenset] = [] def adjudicate(self, a, a_facts, b, b_facts): key = frozenset({a.id, b.id}) self.calls.append(key) return er2.Adjudication(self.verdicts.get(key, "unsure"), f"fake verdict for {sorted(key)}") # ---------------------------------------------------------------------- # 1. Candidate-generation normalization — the exact Johann variants. # NEVER identity: also locked in below. # ---------------------------------------------------------------------- def test_block_normalize_underscore(): assert er2.block_normalize("Johann_Boedecker") == "johann boedecker" def test_block_normalize_diacritic_and_comma(): # diacritic fold (NFD + strip combining) + punctuation strip assert er2.block_normalize("Bödecker, Johann") == "bodecker johann" def test_token_sort_handles_comma_inversion(): assert er2.token_sort("Bödecker, Johann") == er2.token_sort("Bodecker Johann") assert er2.token_sort("Bödecker, Johann") == "bodecker johann" def test_char_trigrams_bridge_oe_vs_diacritic_fold(): # "Boedecker" (oe transliteration) vs "Bödecker" (folds to bodecker) # differ as tokens but must share trigram buckets. shared = er2.char_trigrams("Boedecker") & er2.char_trigrams("Bödecker") assert {"dec", "eck", "cke", "ker"} <= shared def test_blocking_normalization_is_not_identity_normalization(): # The identity scheme (v1 / entity_id.py) preserves diacritics and # punctuation; the blocking form must NEVER replace it. s = "Bödecker, Johann" assert v1._normalize_surface(s) == "bödecker, johann" # identity: untouched ö + comma assert er2.block_normalize(s) == "bodecker johann" # blocking only assert v1._normalize_surface(s) != er2.block_normalize(s) def test_blocking_groups_johann_variants(): a = make_entity("e_a", "Johann_Boedecker") b = make_entity("e_b", "Bödecker, Johann") c = make_entity("e_c", "Johann Boedecker") bare = make_entity("e_bare", "Johann") near = make_entity("e_near", "Johanna Phil") unrelated = make_entity("e_zoe", "Zoe Quist") pairs = er2.generate_candidate_pairs([a, b, c, bare, near, unrelated]) keys = {p.key for p in pairs} # All Boedecker spellings co-block pairwise. assert frozenset({"e_a", "e_b"}) in keys assert frozenset({"e_a", "e_c"}) in keys assert frozenset({"e_b", "e_c"}) in keys # Bare "Johann" blocks with the full names (first/last token key). assert frozenset({"e_bare", "e_c"}) in keys # Near-miss "Johanna Phil" co-blocks via trigrams — it must reach # the similarity/adjudication stage (and be rejected there), not # be silently invisible. assert frozenset({"e_near", "e_c"}) in keys # Unrelated person doesn't pair with the Boedecker cluster. assert not any("e_zoe" in k and k != frozenset({"e_zoe"}) for k in keys) def test_blocking_email_local_part_key(): a = make_entity("e_a", "johann.boedecker@pentatonic.com", aliases=["johann.boedecker@pentatonic.com"]) b = make_entity("e_b", "Johann Boedecker") pairs = er2.generate_candidate_pairs([a, b]) assert frozenset({"e_a", "e_b"}) in {p.key for p in pairs} def test_blocking_skips_pairs_already_grouped_by_v1(): a = make_entity("e_a", "Johann_Boedecker") b = make_entity("e_b", "Johann Boedecker") pairs = er2.generate_candidate_pairs( [a, b], already_grouped={frozenset({"e_a", "e_b"})}) assert pairs == [] def test_oversized_blocks_are_skipped(): clones = [make_entity(f"e_{i}", f"Johann Clone{i}") for i in range(10)] pairs = er2.generate_candidate_pairs(clones, max_block=3) # every key these share ("first:johann"/"last:..."/trigrams of # "johann...") is oversized or unique per clone → no pair explosion # from the shared-first-name block. shared_first = [p for p in pairs if "first:clone0" in p.shared_keys] assert shared_first == [] # ---------------------------------------------------------------------- # 2. Threshold band routing # ---------------------------------------------------------------------- def test_threshold_constants(): assert er2.HIGH_THRESHOLD == 0.92 assert er2.LOW_THRESHOLD == 0.75 def test_route_band_boundaries(): assert er2.route_band(0.95) == "high" assert er2.route_band(0.92) == "high" # inclusive assert er2.route_band(0.9199) == "ambiguous" assert er2.route_band(0.75) == "ambiguous" # inclusive assert er2.route_band(0.7499) == "drop" assert er2.route_band(0.10) == "drop" def _scored_pair(a, b, sim): p = er2.CandidatePair(a=a, b=b) p.similarity = sim return p def test_route_pairs_band_routing(): a1, a2 = make_entity("e_1", "Johann Boedecker"), make_entity("e_2", "Bödecker, Johann") b1, b2 = make_entity("e_3", "Carla Voss"), make_entity("e_4", "Carla Vosse") c1, c2 = make_entity("e_5", "Johanna Phil"), make_entity("e_6", "Johann Boedeker") pairs = [ _scored_pair(a1, a2, 0.96), # high → auto merge, no LLM _scored_pair(b1, b2, 0.85), # ambiguous → adjudicated yes _scored_pair(c1, c2, 0.60), # below low → dropped, no LLM ] adj = FakeAdjudicator({frozenset({"e_3", "e_4"}): "yes"}) routed = er2.route_pairs(pairs, adj, {}) assert {p.key for p in routed.merge} == {frozenset({"e_1", "e_2"}), frozenset({"e_3", "e_4"})} assert {p.key for p in routed.dropped} == {frozenset({"e_5", "e_6"})} assert routed.human_review == [] # LLM consulted ONLY for the ambiguous pair. assert adj.calls == [frozenset({"e_3", "e_4"})] # high-confidence merge carries its reasoning auto = next(p for p in routed.merge if p.key == frozenset({"e_1", "e_2"})) assert auto.verdict == "auto" and "0.960" in auto.reason def test_route_pairs_unsure_never_merges(): a, b = make_entity("e_1", "Johann Boedecker"), make_entity("e_2", "Johanna Phil") routed = er2.route_pairs( [_scored_pair(a, b, 0.85)], FakeAdjudicator({frozenset({"e_1", "e_2"}): "unsure"}), {}) assert routed.merge == [] assert [p.key for p in routed.human_review] == [frozenset({"e_1", "e_2"})] def test_route_pairs_no_verdict_means_no_merge(): a, b = make_entity("e_1", "Johann Boedecker"), make_entity("e_2", "Johanna Phil") routed = er2.route_pairs( [_scored_pair(a, b, 0.85)], FakeAdjudicator({frozenset({"e_1", "e_2"}): "no"}), {}) assert routed.merge == [] and routed.human_review == [] assert [p.key for p in routed.dropped] == [frozenset({"e_1", "e_2"})] def test_no_llm_routes_whole_ambiguous_band_to_human_review(): a, b = make_entity("e_1", "Carla Voss"), make_entity("e_2", "Carla Vosse") c, d = make_entity("e_3", "Mark Diaz"), make_entity("e_4", "Marc Diaz") routed = er2.route_pairs( [_scored_pair(a, b, 0.85), _scored_pair(c, d, 0.80)], er2.NoLLMAdjudicator(), {}) assert routed.merge == [] assert len(routed.human_review) == 2 assert all(p.verdict == "unsure" for p in routed.human_review) # ---------------------------------------------------------------------- # 3. Bare-first-name policy # ---------------------------------------------------------------------- def test_is_bare_name(): assert er2.is_bare_name(make_entity("e", "Johann")) assert not er2.is_bare_name(make_entity("e", "Johann Boedecker")) # an email-only entity is not a bare *name* assert not er2.is_bare_name(make_entity("e", "johann@x.com")) # a single-token canonical with a multi-token alias is not bare assert not er2.is_bare_name( make_entity("e", "Johann", aliases=["Johann Boedecker"])) def test_bare_name_single_candidate_and_yes_merges(): bare = make_entity("e_bare", "Johann") full = make_entity("e_full", "Johann Boedecker", facts=10) routed = er2.route_pairs( [_scored_pair(bare, full, 0.95)], FakeAdjudicator({frozenset({"e_bare", "e_full"}): "yes"}), {}) assert [p.key for p in routed.merge] == [frozenset({"e_bare", "e_full"})] def test_bare_name_high_similarity_still_requires_adjudication(): bare = make_entity("e_bare", "Johann") full = make_entity("e_full", "Johann Boedecker") adj = FakeAdjudicator({frozenset({"e_bare", "e_full"}): "no"}) routed = er2.route_pairs([_scored_pair(bare, full, 0.99)], adj, {}) # 0.99 >= high, but bare names never auto-merge: LLM said no → drop assert routed.merge == [] assert adj.calls == [frozenset({"e_bare", "e_full"})] def test_bare_name_two_candidates_never_merges(): bare = make_entity("e_bare", "Johann") full1 = make_entity("e_f1", "Johann Boedecker") full2 = make_entity("e_f2", "Johann Mueller") adj = FakeAdjudicator({ frozenset({"e_bare", "e_f1"}): "yes", frozenset({"e_bare", "e_f2"}): "yes", }) routed = er2.route_pairs( [_scored_pair(bare, full1, 0.93), _scored_pair(bare, full2, 0.90)], adj, {}) assert routed.merge == [] assert len(routed.human_review) == 2 assert adj.calls == [] # ambiguity short-circuits before the LLM def test_bare_name_with_no_llm_goes_to_review(): bare = make_entity("e_bare", "Johann") full = make_entity("e_full", "Johann Boedecker") routed = er2.route_pairs([_scored_pair(bare, full, 0.95)], er2.NoLLMAdjudicator(), {}) assert routed.merge == [] assert len(routed.human_review) == 1 # ---------------------------------------------------------------------- # 4. Proposal assembly (richest-row-wins, same ordering as v1) # ---------------------------------------------------------------------- def test_pairs_to_proposals_richest_wins_and_transitive(): a = make_entity("e_a", "Johann_Boedecker", facts=2) b = make_entity("e_b", "Bödecker, Johann", facts=50, rels=3) c = make_entity("e_c", "Johann Boedecker", facts=7) p1, p2 = _scored_pair(a, b, 0.95), _scored_pair(b, c, 0.94) proposals = er2.pairs_to_proposals([p1, p2]) assert len(proposals) == 1 assert proposals[0].canonical.id == "e_b" # richest assert {d.id for d in proposals[0].deprecated} == {"e_a", "e_c"} assert proposals[0].signal == "embedding_llm" # ---------------------------------------------------------------------- # 5. Adjudication JSON parsing — strict; malformed → unsure # ---------------------------------------------------------------------- def test_parse_adjudication_valid(): out = er2.parse_adjudication( '{"same_person": "yes", "reason": "same email domain and role"}') assert out.same_person == "yes" assert "email domain" in out.reason def test_parse_adjudication_json_embedded_in_prose(): out = er2.parse_adjudication( 'Sure! Here is my answer: {"same_person": "no", "reason": "different people"}') assert out.same_person == "no" def test_parse_adjudication_malformed_is_unsure(): assert er2.parse_adjudication("I think they are the same").same_person == "unsure" assert er2.parse_adjudication('{"same_person": "maybe"}').same_person == "unsure" assert er2.parse_adjudication('{"broken json').same_person == "unsure" # ---------------------------------------------------------------------- # 6. Embedding plumbing (no network — backends are constructed only) # ---------------------------------------------------------------------- def test_cosine(): assert abs(er2.cosine([1.0, 0.0], [1.0, 0.0]) - 1.0) < 1e-9 assert abs(er2.cosine([1.0, 0.0], [0.0, 1.0])) < 1e-9 assert er2.cosine([0.0, 0.0], [1.0, 0.0]) == 0.0 def test_embedding_bundle_contains_surfaces_and_facts(): e = make_entity("e", "Johann Boedecker", aliases=["johann@x.com"]) bundle = er2.embedding_bundle(e, ["Johann leads the Berlin office"]) assert "Johann Boedecker" in bundle assert "johann@x.com" in bundle assert "Berlin office" in bundle def test_http_backend_requires_url_and_is_never_defaulted(): try: er2.HttpEmbeddingBackend("") raise AssertionError("expected ValueError for empty --embed-url") except ValueError: pass def test_local_backend_is_an_explicit_stub(): try: er2.LocalEmbeddingBackend() raise AssertionError("expected NotImplementedError") except NotImplementedError as e: assert "--embed-backend http" in str(e) # ---------------------------------------------------------------------- # 7. Safety gates: --apply refused without --i-have-a-snapshot # ---------------------------------------------------------------------- def test_validate_args_refuses_apply_without_snapshot(): args = er2.parse_args(["--arena", "test-arena", "--pg-dsn", "x", "--embed-url", "http://e", "--apply"]) err = er2.validate_args(args) assert err is not None and "--i-have-a-snapshot" in err def test_validate_args_accepts_apply_with_snapshot(): args = er2.parse_args(["--arena", "test-arena", "--pg-dsn", "x", "--embed-url", "http://e", "--apply", "--i-have-a-snapshot"]) assert er2.validate_args(args) is None def test_cli_subprocess_refuses_apply_without_snapshot(): proc = subprocess.run( [sys.executable, str(_SCRIPTS / "entity_resolution_v2.py"), "--arena", "test-arena", "--pg-dsn", "postgresql://x", "--apply"], capture_output=True, text=True) assert proc.returncode == 2 assert "--i-have-a-snapshot" in proc.stderr def test_cli_requires_arena(): proc = subprocess.run( [sys.executable, str(_SCRIPTS / "entity_resolution_v2.py"), "--pg-dsn", "postgresql://x"], capture_output=True, text=True) assert proc.returncode == 2 assert "--arena" in proc.stderr # ---------------------------------------------------------------------- # 8. Arena scoping — every SQL statement carries the arena predicate # ---------------------------------------------------------------------- def test_v2_sql_registry_is_arena_scoped(): assert er2.ARENA_SCOPED_SQL, "registry must not be empty" for name, sql in er2.ARENA_SCOPED_SQL.items(): assert "arena = %s" in sql, f"v2 SQL '{name}' lost its arena predicate" er2.assert_arena_scoped() # must not raise def test_v1_load_and_repoint_sql_is_arena_scoped(): src = (_SCRIPTS / "backfill_entity_reconciliation.py").read_text() # Every fact/relationship repoint in apply_proposals is arena-scoped. import re repoints = re.findall( r"UPDATE (?:facts|relationships) SET \w+ = %s\s+WHERE ([^\"]+?)\n", src) assert len(repoints) == 4, "expected the 4 v1 repoint statements" for predicate in repoints: assert "arena = %s" in predicate # The entity load + co-occurrence scan are arena-scoped too. assert "WHERE arena = %s AND entity_type = %s" in src assert '"SELECT attributes FROM events WHERE arena = %s"' in src # ---------------------------------------------------------------------- # 9. v1 CLI byte-compatibility — v2 must not have changed it # ---------------------------------------------------------------------- def test_v1_cli_surface_unchanged(): old_argv = sys.argv try: sys.argv = ["backfill_entity_reconciliation.py", "--arena", "a", "--pg-dsn", "dsn"] ns = v1.parse_args() finally: sys.argv = old_argv assert vars(ns) == { "arena": "a", "pg_dsn": "dsn", "entity_type": "person", "apply": False, "heuristic_merge": False, "out": None, "merged_by": None, }, "v1's CLI namespace changed — it must stay byte-compatible" def test_v1_cli_rejects_v2_only_flags(): old_argv = sys.argv try: sys.argv = ["backfill_entity_reconciliation.py", "--arena", "a", "--pg-dsn", "dsn", "--i-have-a-snapshot"] try: v1.parse_args() raise AssertionError("v1 accepted a v2-only flag") except SystemExit as e: assert e.code == 2 finally: sys.argv = old_argv def test_v1_machinery_is_imported_not_copied(): # v2 must reuse v1's load/apply machinery, not fork it. assert er2.v1 is v1 assert er2.Entity is v1.Entity assert er2.MergeProposal is v1.MergeProposal src = (_SCRIPTS / "entity_resolution_v2.py").read_text() for fn in ("def load_entities", "def apply_proposals", "def collect_cooccurrence_pairs", "def build_proposals"): assert fn not in src, f"v2 copied v1's {fn} instead of importing it" # ---------------------------------------------------------------------- # 10. Tiered report # ---------------------------------------------------------------------- def test_report_tiers_and_counts(tmp_path=None): rich = make_entity("e_rich", "Johann Boedecker", facts=50) dup = make_entity("e_dup", "Johann_Boedecker", facts=1) v1_prop = v1.MergeProposal(canonical=rich, deprecated=[dup], signal="co_occurrence") m1, m2 = make_entity("e_m1", "Carla Voss", facts=9), make_entity("e_m2", "Carla Vosse") pair = _scored_pair(m1, m2, 0.85) pair.band, pair.verdict, pair.reason = "ambiguous", "yes", "same person per facts" routed = er2.RoutedPairs(merge=[pair]) v2_props = er2.pairs_to_proposals([pair]) records = er2.build_report_records( "test-arena", [v1_prop], v2_props, routed, before_count=10) header = records[0] assert header["arena"] == "test-arena" assert header["entity_count_before"] == 10 assert header["entity_count_after_if_applied"] == 8 # 2 deprecations assert "untouched" in header["other_arenas"] assert header["tiers"] == ["co_occurrence", "alias_overlap", "embedding_llm", "heuristic"] proposals = [r for r in records if r["type"] == "merge_proposal"] assert [p["tier"] for p in proposals] == ["co_occurrence", "embedding_llm"] emb = proposals[1] assert emb["evidence"][0]["similarity"] == 0.85 assert emb["evidence"][0]["verdict"] == "yes" assert "same person per facts" in emb["evidence"][0]["reason"] # All records JSONL-serializable. for r in records: json.dumps(r) def test_markdown_summary_mentions_frozen_arena(tmp_path=None): import tempfile, os records = er2.build_report_records( "test-arena", [], [], er2.RoutedPairs(), before_count=3) with tempfile.TemporaryDirectory() as d: path = os.path.join(d, "summary.md") er2.write_markdown_summary(records, path) text = Path(path).read_text() assert "Other arenas" in text assert "pip-agents" in text and "FROZEN" in text assert "`test-arena`" in text