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"