""" Task 1.3.4 — Unit tests for EntityResolver. Registry Lookup is geo-blocked outside the US, so all tests mock _rl_search. EDGAR and Companies House tests include both mocked units and one live integration test (requires internet access; no VPN needed). Run: pytest tests/test_entity_resolver.py -v Run live tests only: pytest tests/test_entity_resolver.py -v -m integration """ import asyncio import json from pathlib import Path from unittest.mock import AsyncMock, patch import pytest from src.models.entities import ResolvedEntity from src.resolution.entity_resolver import ( EntityNotFoundError, EntityResolver, _cache_key, _generate_aliases, _strip_suffixes, _get_cached, _set_cached, _cache_key, ) # ── Utility tests ────────────────────────────────────────────────────────────── class TestHelpers: def test_strip_suffixes_inc(self): assert _strip_suffixes("Tesla, Inc.") == "Tesla" def test_strip_suffixes_llc(self): assert _strip_suffixes("Acme LLC") == "Acme" def test_strip_suffixes_ltd(self): assert _strip_suffixes("Revolut Ltd.") == "Revolut" def test_strip_suffixes_corp(self): assert _strip_suffixes("Boeing Corp") == "Boeing" def test_strip_suffixes_no_suffix(self): assert _strip_suffixes("Stripe") == "Stripe" def test_generate_aliases_deduplicates(self): aliases = _generate_aliases("Tesla, Inc.", "Tesla Inc", ["TSLA"]) assert "Tesla, Inc." in aliases assert "Tesla Inc" in aliases assert "TSLA" in aliases # No duplicates assert len(aliases) == len(set(aliases)) def test_generate_aliases_stripped_variants(self): aliases = _generate_aliases("TESLA, INC.", "Tesla Inc.", ["TSLA"]) # Stripped canonical ("TESLA") and stripped raw ("Tesla Inc") should appear assert "TESLA" in aliases def test_cache_key_stable(self): assert _cache_key("Tesla Inc") == _cache_key("Tesla Inc") def test_cache_key_case_insensitive(self): assert _cache_key("TESLA") == _cache_key("tesla") == _cache_key("Tesla") def test_cache_key_different_inputs(self): assert _cache_key("Tesla") != _cache_key("Stripe") # ── EntityResolver unit tests (all external calls mocked) ──────────────────── MOCK_TESLA_EDGAR = { "cik": "0001318605", "canonical_name": "Tesla, Inc.", "tickers": ["TSLA"], "state": "TX", } MOCK_STRIPE_RL = [ { "id": "abc123", "legal_name": "Stripe, Inc.", "jurisdiction_code": "us-de", "company_type": "Private", } ] MOCK_REVOLUT_RL = [ { "id": "uk001", "legal_name": "REVOLUT LTD", "jurisdiction_code": "gb", "company_type": "Private Limited Company", } ] MOCK_REVOLUT_CH = [ { "title": "REVOLUT LTD", "company_number": "08804411", "company_status": "active", } ] MOCK_AMBIGUOUS_RL = [ {"id": "1", "legal_name": "Mercury Corp A", "jurisdiction_code": "us-de"}, {"id": "2", "legal_name": "Mercury Corp B", "jurisdiction_code": "us-ca"}, {"id": "3", "legal_name": "Mercury Corp C", "jurisdiction_code": "us-ny"}, {"id": "4", "legal_name": "Mercury Corp D", "jurisdiction_code": "gb"}, ] @pytest.fixture def resolver(): return EntityResolver() class TestEntityResolverUnit: """All Registry Lookup + EDGAR + Companies House calls are mocked.""" async def test_resolve_known_public_company(self, resolver, tmp_path, monkeypatch): """Tesla: RL geo-blocked (returns []), EDGAR returns CIK → is_public=True.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=([], True)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value=MOCK_TESLA_EDGAR): entity = await resolver.resolve("Tesla Inc") assert entity.canonical_name == "Tesla, Inc." assert entity.is_public is True assert entity.sec_cik == "0001318605" assert "TSLA" in entity.aliases assert "Tesla" in entity.aliases or "Tesla Inc" in entity.aliases async def test_resolve_known_private_company(self, resolver, tmp_path, monkeypatch): """Stripe: RL returns result, EDGAR returns nothing → is_public=False.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(MOCK_STRIPE_RL, False)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value={}): entity = await resolver.resolve("Stripe") assert "Stripe" in entity.canonical_name assert entity.is_public is False assert entity.sec_cik is None assert entity.jurisdiction == "us-de" async def test_resolve_uk_company_enriched(self, resolver, tmp_path, monkeypatch): """Revolut: RL returns gb jurisdiction → Companies House enriched.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(MOCK_REVOLUT_RL, False)), \ patch("src.resolution.entity_resolver._ch_search", new_callable=AsyncMock, return_value=MOCK_REVOLUT_CH), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value={}): entity = await resolver.resolve("Revolut") assert entity.companies_house_number == "08804411" assert "REVOLUT" in entity.canonical_name.upper() async def test_resolve_ambiguous_name(self, resolver, tmp_path, monkeypatch): """>3 RL matches → auto-select the top-ranked match (never dead-end).""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(MOCK_AMBIGUOUS_RL, False)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value={}): entity = await resolver.resolve("Mercury Corp") # The top match is selected; the pipeline proceeds instead of erroring. assert isinstance(entity, ResolvedEntity) assert entity.canonical_name == "Mercury Corp A" async def test_resolve_not_found(self, resolver, tmp_path, monkeypatch): """RL empty, EDGAR empty → EntityNotFoundError.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=([], False)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value={}), \ pytest.raises(EntityNotFoundError): await resolver.resolve("XYZ_DOES_NOT_EXIST_12345") async def test_cache_hit(self, resolver, tmp_path, monkeypatch): """Second resolve call uses cache — external APIs not called again.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", True) rl_mock = AsyncMock(return_value=(MOCK_STRIPE_RL, False)) edgar_mock = AsyncMock(return_value={}) with patch("src.resolution.entity_resolver._rl_search", rl_mock), \ patch("src.resolution.entity_resolver._edgar_lookup", edgar_mock): e1 = await resolver.resolve("Stripe") e2 = await resolver.resolve("Stripe") assert e1.canonical_name == e2.canonical_name # RL should have been called only once (second call served from cache) assert rl_mock.call_count == 1 async def test_edgar_fallback_when_rl_blocked(self, resolver, tmp_path, monkeypatch): """RL returns empty (geo-blocked), EDGAR finds the company.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=([], True)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value=MOCK_TESLA_EDGAR): entity = await resolver.resolve("Tesla") assert entity.is_public is True assert entity.sec_cik == "0001318605" async def test_ambiguity_selects_first_of_many(self, resolver, tmp_path, monkeypatch): """Many matches → auto-select the first-ranked, never raise.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) many = [{"id": str(i), "legal_name": f"Acme #{i}", "jurisdiction_code": "us-de"} for i in range(10)] with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(many, False)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value={}): entity = await resolver.resolve("Acme") assert entity.canonical_name == "Acme #0" async def test_search_candidates_returns_display_views(self, resolver, tmp_path, monkeypatch): """search_candidates returns compact, display-ready dicts (picker §8c).""" monkeypatch.setattr("src.config.settings.use_cache", False) results = [ {"id": "1", "legal_name": "Formlabs Inc", "jurisdiction_code": "us-ma", "status": "active", "company_type": "corporation", # Registry Lookup often returns address as a nested object, not a string. "registered_address": {"locality": "Somerville", "region": "MA", "country": "US"}}, {"id": "2", "legal_name": "Formlabs KK", "jurisdiction_code": "jp", "status": "active"}, ] with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(results, False)): cands = await resolver.search_candidates("Formlabs") assert [c["registry_id"] for c in cands] == ["1", "2"] assert cands[0]["name"] == "Formlabs Inc" and cands[0]["jurisdiction"] == "us-ma" assert cands[0]["status"] == "active" assert cands[1]["jurisdiction"] == "jp" # Nested address object must be flattened to a string (else /api/resolve 500s). assert isinstance(cands[0]["address"], str) and "Somerville" in cands[0]["address"] async def test_resolve_by_registry_id_selects_exact(self, resolver, tmp_path, monkeypatch): """resolve(registry_id=...) picks that exact candidate, not the top-ranked.""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) # Top-ranked is the JP entity; the user picked the US one (id=2). results = [ {"id": "1", "legal_name": "Formlabs KK", "jurisdiction_code": "jp"}, {"id": "2", "legal_name": "Formlabs Inc", "jurisdiction_code": "us-ma"}, ] with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(results, False)), \ patch("src.resolution.entity_resolver._edgar_lookup", new_callable=AsyncMock, return_value={}): entity = await resolver.resolve("Formlabs", registry_id="2") assert entity.canonical_name == "Formlabs Inc" assert entity.registry_lookup_id == "2" # ── Integration tests (live API — no VPN needed) ────────────────────────────── @pytest.mark.integration class TestEntityResolverIntegration: """These tests call real APIs. Registry Lookup will return [] (geo-blocked) but EDGAR and Companies House are accessible from any country.""" async def test_resolve_tesla_live(self, tmp_path, monkeypatch): """Tesla resolves via EDGAR (RL geo-blocked from India).""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) resolver = EntityResolver() with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=([], True)): entity = await resolver.resolve("Tesla") assert entity.is_public is True assert "1318605" in entity.sec_cik assert "TSLA" in entity.aliases assert "Tesla" in entity.canonical_name async def test_resolve_stripe_via_rl_mock_and_edgar(self, tmp_path, monkeypatch): """Stripe: RL mocked (geo-blocked), EDGAR should return no CIK (private).""" monkeypatch.setattr("src.config.settings.database_path", str(tmp_path / "test.db")) monkeypatch.setattr("src.config.settings.use_cache", False) resolver = EntityResolver() with patch("src.resolution.entity_resolver._rl_search", new_callable=AsyncMock, return_value=(MOCK_STRIPE_RL, False)): entity = await resolver.resolve("Stripe") assert entity.is_public is False assert entity.sec_cik is None