tests/test_embedder.py

from unittest.mock import patch
from ingestion.embedder import chunk_text, embed_and_store_filing, embed_and_store_transcript


def test_chunk_long_text_produces_multiple_chunks():
    text = "word " * 600
    chunks = chunk_text(text, chunk_size=100, overlap=10)
    assert len(chunks) > 1


def test_chunk_short_text_returns_single_chunk():
    assert chunk_text("Short text.") == ["Short text."]


def test_chunk_overlap_repeats_words():
    text = " ".join(str(i) for i in range(200))
    chunks = chunk_text(text, chunk_size=100, overlap=20)
    last_word = chunks[0].split()[-1]
    assert last_word in chunks[1]


@patch("ingestion.embedder.vector_store")
def test_embed_filing_calls_add_chunks_twice(mock_vs):
    embed_and_store_filing(
        ticker="AAPL",
        company_name="Apple Inc.",
        mda_text="Revenue grew 5%.",
        risk_text="Competition remains a risk.",
        filing_date="2024-11-01",
        period="Q42024",
        form_type="10-Q",
    )
    assert mock_vs.add_chunks.call_count == 2


@patch("ingestion.embedder.vector_store")
def test_embed_filing_skips_empty_sections(mock_vs):
    embed_and_store_filing(
        ticker="AAPL",
        company_name="Apple Inc.",
        mda_text="Some MD&A text.",
        risk_text="",
        filing_date="2024-11-01",
        period="Q42024",
        form_type="10-Q",
    )
    assert mock_vs.add_chunks.call_count == 1


@patch("ingestion.embedder.vector_store")
def test_embed_annual_filing_indexes_company_profile_sections(mock_vs):
    embed_and_store_filing(
        ticker="AAPL",
        company_name="Apple Inc.",
        mda_text="Revenue grew 5%.",
        risk_text="Competition remains a risk.",
        business_text="The company designs and sells devices and services.",
        segments_geography_text="Net sales are reported by geography.",
        filing_date="2024-11-01",
        period="FY2024",
        form_type="10-K",
    )

    assert mock_vs.add_chunks.call_count == 4
    metadata = [call.args[2][0] for call in mock_vs.add_chunks.call_args_list]
    assert [item["section"] for item in metadata] == [
        "MD&A",
        "Risk Factors",
        "Business",
        "Segments & Geography",
    ]
    assert [item["chunk_id"] for item in metadata] == [
        "mda:0",
        "risk_factors:0",
        "business:0",
        "segments_geography:0",
    ]


@patch("ingestion.embedder.vector_store")
def test_embed_filing_chunk_context_in_metadata(mock_vs):
    embed_and_store_filing(
        ticker="AAPL",
        company_name="Apple Inc.",
        mda_text="Revenue grew 5%.",
        risk_text="",
        filing_date="2024-11-01",
        period="Q42024",
        form_type="10-Q",
    )
    _, kwargs = mock_vs.add_chunks.call_args_list[0]
    metadatas = mock_vs.add_chunks.call_args_list[0][0][2]
    assert "chunk_context" in metadatas[0]
    assert "Apple Inc." in metadatas[0]["chunk_context"]
    assert "Q42024" in metadatas[0]["chunk_context"]
    assert metadatas[0]["document_id"].startswith("sec:AAPL:")
    assert metadatas[0]["chunk_id"] == "mda:0"


@patch("ingestion.embedder.vector_store")
def test_embed_transcript(mock_vs):
    embed_and_store_transcript(
        ticker="AAPL",
        company_name="Apple Inc.",
        transcript_text="Tim Cook: Record quarter.",
        transcript_date="2024-11-01",
        period="Q42024",
    )
    mock_vs.add_chunks.assert_called_once()
    args = mock_vs.add_chunks.call_args
    assert args[0][0] == "transcripts"
    metadatas = args[0][2]
    assert "chunk_context" in metadatas[0]
    assert "Q42024" in metadatas[0]["chunk_context"]
    assert metadatas[0]["document_id"] == "transcript:AAPL:Q42024:2024-11-01"
    assert metadatas[0]["chunk_id"] == "transcript:0"