CIAgent/apps/api/tests/unit/test_enrichment_service.py

"""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