File size: 3,754 Bytes
35676b4 d1e793b 35676b4 7880373 35676b4 7880373 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | 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"
|