Files
CIAgent/apps/api/tests/unit/test_enrichment_service.py
sakshamandClaude Sonnet 5 18305b545c Backfill sparse report sections, add enrichment section refresh, and bootstrap first-admin
Reports: the LLM reliably used company_enrichment for prose fields but
inconsistently populated the parallel Finding-list/string-list fields from
the same evidence, even with progressively more explicit prompting. Add a
code-level backfill (products, recent developments, financial signals,
strategic initiatives, regulatory signals, risks/opportunities mirrored
from SWOT, unknowns, monitoring recommendations) that only ever fills in
what the model left empty, never overwrites what it produced.

Enrichment tab: reorder sections (Products/Recent updates before
Customers/Competitors) and add a per-section "Refresh" button that
re-fetches just one of NinjaPear's six independent per-company endpoints
when it came back empty - confirmed live that a data-coverage gap (e.g.
Amazon returning no products) is real provider behavior, not a bug.

Auth: the first account registered on a deployment with zero existing
admins is now auto-promoted to admin, closing the chicken-and-egg gap
where the only path to admin access was direct DB access. Self-heals if
the last admin ever deletes their account.

Also bumps nginx's proxy_read_timeout for api.ciagent.org to cover the
enrichment refresh's synchronous funding-endpoint call (up to 5 minutes
per NinjaPear's docs).

Co-Authored-By: Claude Sonnet 5 <[email protected]>
2026-08-06 17:41:02 -04:00

267 lines
10 KiB
Python

"""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.core.errors import NotFoundError
from app.enrichment.base import (
CompanyDetails,
CompanyFunding,
LeadershipMember,
Product,
)
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, refresh_section
class _FakeProvider:
provider_name = "fake"
def __init__(
self,
*,
leadership: list[LeadershipMember] | None = None,
products: list[Product] | None = None,
fail: set[str] | None = None,
):
self._leadership = leadership or []
self._products = products if products is not None else []
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 self._products
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
@pytest.mark.asyncio
async def test_refresh_section_replaces_only_that_section(db_session, settings):
company = await _make_company(db_session)
await enrich_company(db_session, settings, _FakeProvider(), company) # products come back []
provider = _FakeProvider(products=[Product(name="Acme Pay", description="Payments product")])
enrichment = await refresh_section(db_session, settings, provider, company, "products")
assert [p["name"] for p in enrichment.data["products"]] == ["Acme Pay"]
assert enrichment.data["description"] == "A company." # untouched by the refresh
assert enrichment.status == EnrichmentStatus.COMPLETE
@pytest.mark.asyncio
async def test_refresh_section_records_a_fresh_failure_without_losing_other_data(
db_session, settings
):
company = await _make_company(db_session)
await enrich_company(db_session, settings, _FakeProvider(), company)
failing_provider = _FakeProvider(fail={"products"})
enrichment = await refresh_section(db_session, settings, failing_provider, company, "products")
assert enrichment.errors["products"] == "products boom"
assert enrichment.data["description"] == "A company."
assert enrichment.status == EnrichmentStatus.PARTIAL
@pytest.mark.asyncio
async def test_refresh_section_clears_a_previously_recorded_error(db_session, settings):
company = await _make_company(db_session)
await enrich_company(db_session, settings, _FakeProvider(fail={"products"}), company)
provider = _FakeProvider(products=[Product(name="Acme Pay")])
enrichment = await refresh_section(db_session, settings, provider, company, "products")
assert "products" not in enrichment.errors
assert enrichment.status == EnrichmentStatus.COMPLETE
@pytest.mark.asyncio
async def test_refresh_section_preserves_per_leader_lookups_on_a_details_refresh(
db_session, settings
):
company = await _make_company(db_session)
leadership = [LeadershipMember(name="Jane Doe")]
await enrich_company(db_session, settings, _FakeProvider(leadership=leadership), company)
stored = await CompanyEnrichmentRepository(db_session).get_for_company(company.id)
assert stored.data["leadership_team"][0]["work_email"] == "[email protected]"
# A details refresh re-fetches the roster but not the per-leader lookups -
# the previously-found work_email must survive, not silently disappear.
provider = _FakeProvider(leadership=[LeadershipMember(name="Jane Doe", title="CEO")])
enrichment = await refresh_section(db_session, settings, provider, company, "details")
member = enrichment.data["leadership_team"][0]
assert member["title"] == "CEO"
assert member["work_email"] == "[email protected]"
@pytest.mark.asyncio
async def test_refresh_section_rejects_an_unknown_section(db_session, settings):
company = await _make_company(db_session)
await enrich_company(db_session, settings, _FakeProvider(), company)
with pytest.raises(ValueError):
await refresh_section(db_session, settings, _FakeProvider(), company, "not_a_section")
@pytest.mark.asyncio
async def test_refresh_section_requires_an_existing_enrichment_record(db_session, settings):
company = await _make_company(db_session) # enrichment was never run
with pytest.raises(NotFoundError):
await refresh_section(db_session, settings, _FakeProvider(), company, "products")
@pytest.mark.asyncio
async def test_refresh_section_adds_its_own_credit_cost(db_session, settings):
company = await _make_company(db_session)
baseline = await enrich_company(db_session, settings, _FakeProvider(), company)
baseline_credits = baseline.credits_spent # `upsert` mutates rows in place - snapshot first
enrichment = await refresh_section(db_session, settings, _FakeProvider(), company, "products")
assert enrichment.credits_spent == baseline_credits + 3 # products' own credit cost