tests/test_coverage.py

from pathlib import Path

from financial_search_engine.config import Settings
from financial_search_engine.schemas import (
    CoverageIngestionRequest,
    CoverageStatus,
    ExpansionUniverse,
    ExtractedConstraints,
    SearchRequest,
)
from financial_search_engine.services.coverage import CoverageService


class FakeStore:
    def __init__(self, coverage=None, years=None, indexed_companies=None):
        self.coverage = coverage or {}
        self.years = years or {}
        self.indexed_companies = indexed_companies or []
        self.created_jobs = []

    def ensure_index(self):
        return None

    def ensure_supporting_indices(self):
        return None

    def ensure_search_pipeline(self):
        return None

    def list_indexed_companies(self):
        return self.indexed_companies

    def find_coverage_record(self, *, ticker=None, company_name=None, fiscal_year=None):
        key = (ticker or company_name, fiscal_year)
        return self.coverage.get(key)

    def get_available_years(self, *, ticker=None, company_name=None):
        key = ticker or company_name
        return self.years.get(key, [])

    def find_active_job(self, request_key):
        return None

    def create_ingestion_job(self, job):
        self.created_jobs.append(job)


def _service(store, settings=None):
    return CoverageService(
        settings=settings or Settings(),
        store=store,
        embedder_factory=lambda: None,
    )


def test_assess_search_coverage_reports_missing_company_year():
    service = _service(FakeStore(years={"TSLA": [2024]}))
    request = SearchRequest(query="Did TSLA mention battery supply risks in 2023?")
    constraints = ExtractedConstraints(tickers=["TSLA"], temporal_years=[2023])

    assessment = service.assess_search_coverage(request, constraints)

    assert assessment.status == CoverageStatus.MISSING
    assert len(assessment.missing_targets) == 1
    assert assessment.available_years_by_identifier["TSLA"] == [2024]


def test_assess_search_coverage_uses_current_indexed_universe_for_broad_queries():
    store = FakeStore(
        coverage={("NVDA", 2023): {"ticker": "NVDA", "fiscal_year": 2023}},
        years={"NVDA": [2023], "AAPL": [2021]},
        indexed_companies=[
            {"ticker": "NVDA", "company_name": "NVIDIA Corporation", "sector_gics": "Information Technology"},
            {"ticker": "AAPL", "company_name": "Apple Inc.", "sector_gics": "Information Technology"},
        ],
    )
    service = _service(store)
    request = SearchRequest(query="Which semiconductor companies flagged supply chain risks in Asia in 2023?")
    constraints = ExtractedConstraints(industrial_sectors=["semiconductors"], temporal_years=[2023])

    assessment = service.assess_search_coverage(request, constraints)

    assert assessment.status == CoverageStatus.PARTIAL
    assert assessment.scope_label == "currently indexed semiconductors universe"
    assert len(assessment.requested_targets) == 2
    assert len(assessment.available_targets) == 1
    assert len(assessment.missing_targets) == 1


def test_assess_search_coverage_matches_specific_indexed_sector_labels():
    store = FakeStore(
        coverage={("NVDA", 2023): {"ticker": "NVDA", "fiscal_year": 2023}},
        years={"NVDA": [2023], "MSFT": [2023]},
        indexed_companies=[
            {"ticker": "NVDA", "company_name": "NVIDIA Corporation", "sector_gics": "Semiconductors & Related Devices"},
            {"ticker": "MSFT", "company_name": "Microsoft Corporation", "sector_gics": "Services-Prepackaged Software"},
        ],
    )
    service = _service(store)
    request = SearchRequest(query="Which semiconductor companies flagged supply chain risks in Asia in 2023?")
    constraints = ExtractedConstraints(industrial_sectors=["semiconductors"], temporal_years=[2023])

    assessment = service.assess_search_coverage(request, constraints)

    assert assessment.status == CoverageStatus.READY
    assert len(assessment.requested_targets) == 1
    assert assessment.requested_targets[0].ticker == "NVDA"


def test_queue_ingestion_resolves_sp100_template_targets(tmp_path: Path):
    csv_path = tmp_path / "sp100.csv"
    csv_path.write_text(
        "ticker,company_name,sector_gics\nNVDA,NVIDIA Corporation,Information Technology\nMSFT,Microsoft Corporation,Information Technology\n",
        encoding="utf-8",
    )
    settings = Settings(SP100_UNIVERSE_CSV_PATH=str(csv_path))
    store = FakeStore()
    service = _service(store, settings=settings)

    response = service.queue_ingestion(
        CoverageIngestionRequest(
            query="Which companies discussed AI infrastructure in FY2022?",
            expansion_universe=ExpansionUniverse.SP100_TEMPLATE,
            requested_years=[2022],
        )
    )

    assert len(response.jobs) == 2
    assert {job.ticker for job in response.jobs} == {"NVDA", "MSFT"}


def test_queue_ingestion_skips_leading_blank_lines_in_reference_csv(tmp_path: Path):
    csv_path = tmp_path / "sp100.csv"
    csv_path.write_text(
        "\n\nticker,company_name,sector_gics\nNVDA,NVIDIA Corporation,Information Technology\n",
        encoding="utf-8",
    )
    settings = Settings(SP100_UNIVERSE_CSV_PATH=str(csv_path))
    store = FakeStore()
    service = _service(store, settings=settings)

    response = service.queue_ingestion(
        CoverageIngestionRequest(
            query="Which semiconductor companies flagged supply chain risks in Asia in FY2021?",
            expansion_universe=ExpansionUniverse.SP100_TEMPLATE,
            requested_years=[2021],
        )
    )

    assert len(response.jobs) == 1
    assert response.jobs[0].ticker == "NVDA"