saksham
1a4c80958f
FastAPI + Celery + Next.js + Postgres/Redis app with company monitoring, source collection, LLM-based change analysis, enrichment, and account security (Turnstile, escalating lockout, email verification).
106 lines
3.5 KiB
Python
106 lines
3.5 KiB
Python
"""app.tasks.enrichment.enrich_company: the Celery task end-to-end
|
|
(mocked HTTP, real DB) - confirms the CompanyEnrichment row persists with
|
|
the expected status after the task runs, and that a missing company is a
|
|
clean no-op rather than a crash."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
|
|
from app.core.config import Settings
|
|
from app.models.company import Company
|
|
from app.repositories.company_enrichment_repository import CompanyEnrichmentRepository
|
|
from app.tasks.enrichment import _enrich_company_async
|
|
|
|
|
|
async def _make_company(db_session) -> Company:
|
|
company = Company(
|
|
id=uuid.uuid4(),
|
|
user_id=uuid.uuid4(),
|
|
name="Acme Corp",
|
|
slug=f"acme-{uuid.uuid4().hex[:6]}",
|
|
official_website="https://acme.example.com",
|
|
)
|
|
db_session.add(company)
|
|
await db_session.commit()
|
|
return company
|
|
|
|
|
|
def _mock_empty_endpoints() -> None:
|
|
for path in (
|
|
"company/details",
|
|
"company/funding",
|
|
"company/updates",
|
|
"competitor/listing",
|
|
"product/listing",
|
|
"customer/listing",
|
|
):
|
|
respx.get(f"https://nubela.co/api/v1/{path}").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enrich_company_task_persists_a_complete_result(db_session, monkeypatch):
|
|
monkeypatch.setattr(
|
|
"app.tasks.enrichment.get_settings", lambda: Settings(ninjapear_api_key="test-key")
|
|
)
|
|
company = await _make_company(db_session)
|
|
|
|
with respx.mock:
|
|
_mock_empty_endpoints()
|
|
await _enrich_company_async(str(company.id))
|
|
|
|
enrichment = await CompanyEnrichmentRepository(db_session).get_for_company(company.id)
|
|
assert enrichment is not None
|
|
assert enrichment.status.value == "complete"
|
|
assert enrichment.fetched_at is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enrich_company_task_records_partial_status_on_a_failed_section(
|
|
db_session, monkeypatch
|
|
):
|
|
monkeypatch.setattr(
|
|
"app.tasks.enrichment.get_settings", lambda: Settings(ninjapear_api_key="test-key")
|
|
)
|
|
company = await _make_company(db_session)
|
|
|
|
with respx.mock:
|
|
respx.get("https://nubela.co/api/v1/company/details").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
respx.get("https://nubela.co/api/v1/company/funding").mock(
|
|
return_value=httpx.Response(500, text="boom")
|
|
)
|
|
respx.get("https://nubela.co/api/v1/company/updates").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
respx.get("https://nubela.co/api/v1/competitor/listing").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
respx.get("https://nubela.co/api/v1/product/listing").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
respx.get("https://nubela.co/api/v1/customer/listing").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
await _enrich_company_async(str(company.id))
|
|
|
|
enrichment = await CompanyEnrichmentRepository(db_session).get_for_company(company.id)
|
|
assert enrichment.status.value == "partial"
|
|
assert "funding" in enrichment.errors
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enrich_company_task_is_a_no_op_for_a_missing_company(db_session, monkeypatch):
|
|
monkeypatch.setattr(
|
|
"app.tasks.enrichment.get_settings", lambda: Settings(ninjapear_api_key="test-key")
|
|
)
|
|
# Should return cleanly rather than raising - no company, nothing to do.
|
|
await _enrich_company_async(str(uuid.uuid4()))
|