"""enrichment_service.enrich_company: per-section independent failure handling, the leadership-lookup cap, and overall status derivation.""" from __future__ import annotations import uuid import pytest from app.enrichment.base import ( CompanyDetails, CompanyFunding, LeadershipMember, ) from app.models.company import Company from app.models.enums import EnrichmentStatus from app.repositories.company_enrichment_repository import CompanyEnrichmentRepository from app.services.enrichment_service import enrich_company class _FakeProvider: provider_name = "fake" def __init__( self, *, leadership: list[LeadershipMember] | None = None, fail: set[str] | None = None ): self._leadership = leadership or [] self._fail = fail or set() async def get_company_details(self, name, website): if "details" in self._fail: raise RuntimeError("details boom") return CompanyDetails(description="A company.", leadership_team=self._leadership) async def get_funding(self, name, website): if "funding" in self._fail: raise RuntimeError("funding boom") return CompanyFunding(total_raised="$1M") async def get_updates(self, name, website): if "updates" in self._fail: raise RuntimeError("updates boom") return [] async def get_competitors(self, name, website): if "competitors" in self._fail: raise RuntimeError("competitors boom") return [] async def get_products(self, name, website): if "products" in self._fail: raise RuntimeError("products boom") return [] async def get_customers(self, name, website): if "customers" in self._fail: raise RuntimeError("customers boom") return [] async def get_work_email(self, person_name, company_website): if "work_email" in self._fail: raise RuntimeError("email boom") return f"{person_name.split()[0].lower()}@example.com" async def get_person_profile(self, person_name, company_website): if "person_profile" in self._fail: raise RuntimeError("profile boom") return f"https://example.com/{person_name}", "A bio." async def _make_company(db_session, *, website: str | None = "https://acme.example.com") -> Company: company = Company( id=uuid.uuid4(), user_id=uuid.uuid4(), name="Acme Corp", slug=f"acme-{uuid.uuid4().hex[:6]}", official_website=website, ) db_session.add(company) await db_session.commit() return company @pytest.mark.asyncio async def test_all_sections_succeeding_yields_complete_status(db_session, settings): company = await _make_company(db_session) provider = _FakeProvider() enrichment = await enrich_company(db_session, settings, provider, company) assert enrichment.status == EnrichmentStatus.COMPLETE assert enrichment.errors == {} assert enrichment.data["funding"]["total_raised"] == "$1M" assert enrichment.credits_spent is not None and enrichment.credits_spent > 0 @pytest.mark.asyncio async def test_one_failed_section_yields_partial_without_losing_the_rest(db_session, settings): company = await _make_company(db_session) provider = _FakeProvider(fail={"funding"}) enrichment = await enrich_company(db_session, settings, provider, company) assert enrichment.status == EnrichmentStatus.PARTIAL assert "funding" in enrichment.errors assert "funding" not in enrichment.data assert enrichment.data["description"] == "A company." # the other sections still ran @pytest.mark.asyncio async def test_every_section_failing_yields_failed_status(db_session, settings): company = await _make_company(db_session) provider = _FakeProvider( fail={"details", "funding", "updates", "competitors", "products", "customers"} ) enrichment = await enrich_company(db_session, settings, provider, company) assert enrichment.status == EnrichmentStatus.FAILED assert enrichment.data.get("leadership_team") == [] @pytest.mark.asyncio async def test_leadership_lookups_are_capped(db_session, settings): company = await _make_company(db_session) leadership = [LeadershipMember(name=f"Person {i}") for i in range(8)] provider = _FakeProvider(leadership=leadership) capped_settings = settings.model_copy(update={"ninjapear_max_leadership_lookups": 3}) enrichment = await enrich_company(db_session, capped_settings, provider, company) team = enrichment.data["leadership_team"] assert len(team) == 8 # every discovered leader is kept... with_email = [m for m in team if m.get("work_email")] assert len(with_email) == 3 # ...but only the first 3 get person-level lookups @pytest.mark.asyncio async def test_no_website_fails_immediately_without_calling_the_provider(db_session, settings): """NinjaPear identifies a company by website only - every call would fail identically, so this must short-circuit to FAILED before spending any credits, rather than attempting (and paying for) doomed calls.""" company = await _make_company(db_session, website=None) class _ExplodingProvider: provider_name = "exploding" async def get_company_details(self, *a, **k): raise AssertionError("must not be called without a website") enrichment = await enrich_company(db_session, settings, _ExplodingProvider(), company) assert enrichment.status == EnrichmentStatus.FAILED assert enrichment.credits_spent == 0 assert "details" in enrichment.errors @pytest.mark.asyncio async def test_upsert_updates_the_same_row_on_a_second_call(db_session, settings): company = await _make_company(db_session) first = await enrich_company(db_session, settings, _FakeProvider(), company) second = await enrich_company(db_session, settings, _FakeProvider(fail={"funding"}), company) assert first.id == second.id assert second.status == EnrichmentStatus.PARTIAL stored = await CompanyEnrichmentRepository(db_session).get_for_company(company.id) assert stored.id == first.id assert stored.status == EnrichmentStatus.PARTIAL