"""Unit tests for embeddings module.""" import pytest from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage from src.core.exceptions import IndexingError from src.extraction.models import ExtractedRecord from src.knowledge.embeddings import EmbeddingService, record_to_text class TestRecordToText: """Tests for record_to_text function.""" def test_simple_record(self): """Test text representation of simple record.""" record = ExtractedRecord( content={"Item Code": "2024-4_0019", "Report": "not applicable"}, headers=["Item Code", "Report"], _source={ "document_id": str(uuid4()), "filename": "distribution_spec.xlsx", "sheet": "Item Mapping", "table_id": 1, "row": 42, "col_range": "A:D", }, ) text = record_to_text(record) assert "Item Code: 2024-4_0019" in text assert "Report: not applicable" in text assert "From: distribution_spec.xlsx" in text assert "Sheet: Item Mapping" in text assert "Row: 42" in text def test_record_with_resolved_content(self): """Test text representation includes resolved symbol meanings.""" record = ExtractedRecord( content={"Item Code": "2024-4_0019", "Screen": "○"}, headers=["Item Code", "Screen"], _source={ "document_id": str(uuid4()), "filename": "distribution_spec.xlsx", "sheet": "Item Mapping", "table_id": 1, "row": 42, "col_range": "A:D", }, ) # Add resolved_content attribute (dynamically added by symbol resolution) record.resolved_content = { "Item Code": "2024-4_0019", "Screen": "○", "Screen_resolved": "Applicable/Yes", } text = record_to_text(record) assert "Item Code: 2024-4_0019" in text assert "Screen: ○" in text assert "Screen_resolved: Applicable/Yes" in text assert "From: distribution_spec.xlsx" in text def test_record_with_none_values(self): """Test that None values are skipped in text representation.""" record = ExtractedRecord( content={"Item Code": "2024-4_0019", "Report": None, "Status": "Active"}, headers=["Item Code", "Report", "Status"], _source={ "document_id": str(uuid4()), "filename": "test.xlsx", "sheet": "Sheet1", "table_id": 1, "row": 10, "col_range": "A:C", }, ) text = record_to_text(record) assert "Item Code: 2024-4_0019" in text assert "Status: Active" in text assert "Report: None" not in text # None values should be skipped assert "From: test.xlsx" in text def test_record_with_special_characters(self): """Test text representation handles special characters.""" record = ExtractedRecord( content={"Name": "Test & Co.", "Symbol": "○/×", "Value": "50%"}, headers=["Name", "Symbol", "Value"], _source={ "document_id": str(uuid4()), "filename": "special_chars.xlsx", "sheet": "Sheet1", "table_id": 1, "row": 5, "col_range": "A:C", }, ) text = record_to_text(record) assert "Name: Test & Co." in text assert "Symbol: ○/×" in text assert "Value: 50%" in text def test_empty_content_edge_case(self): """Test handling of record with empty content dict.""" record = ExtractedRecord( content={"Empty": None}, headers=["Empty"], _source={ "document_id": str(uuid4()), "filename": "test.xlsx", "sheet": "Sheet1", "table_id": 1, "row": 1, "col_range": "A:A", }, ) text = record_to_text(record) # Should at least include source context assert "From: test.xlsx" in text assert "Sheet: Sheet1" in text assert "Row: 1" in text class TestEmbeddingService: """Tests for EmbeddingService class.""" @pytest.fixture def mock_openai_client(self): """Create mock Azure OpenAI client.""" client = MagicMock() client.embeddings = MagicMock() return client @pytest.fixture def embedding_service(self, mock_openai_client): """Create EmbeddingService with mocked client.""" with patch("src.knowledge.embeddings.AsyncAzureOpenAI") as mock_class: mock_class.return_value = mock_openai_client service = EmbeddingService() service.client = mock_openai_client return service @pytest.mark.asyncio async def test_embed_batch_success(self, embedding_service, mock_openai_client): """Test successful batch embedding.""" texts = [ "Item Code: 2024-4_0019. Report: not applicable", "Item Code: 2024-4_0020. Report: applicable", ] # Mock response with 3072-dimensional vectors mock_embedding_1 = MagicMock(spec=Embedding) mock_embedding_1.embedding = [0.1] * 3072 mock_embedding_2 = MagicMock(spec=Embedding) mock_embedding_2.embedding = [0.2] * 3072 mock_response = MagicMock(spec=CreateEmbeddingResponse) mock_response.data = [mock_embedding_1, mock_embedding_2] mock_openai_client.embeddings.create = AsyncMock(return_value=mock_response) # Call embed_batch embeddings = await embedding_service.embed_batch(texts) # Verify results assert len(embeddings) == 2 assert len(embeddings[0]) == 3072 assert len(embeddings[1]) == 3072 assert embeddings[0] == [0.1] * 3072 assert embeddings[1] == [0.2] * 3072 # Verify API call mock_openai_client.embeddings.create.assert_called_once() call_args = mock_openai_client.embeddings.create.call_args assert call_args.kwargs["input"] == texts assert call_args.kwargs["dimensions"] == 3072 @pytest.mark.asyncio async def test_embed_batch_exceeds_limit(self, embedding_service): """Test error when batch size exceeds 100.""" texts = [f"Text {i}" for i in range(101)] with pytest.raises(ValueError, match="Batch size must be ≤ 100"): await embedding_service.embed_batch(texts) @pytest.mark.asyncio async def test_embed_batch_empty_list(self, embedding_service): """Test error when batch is empty.""" with pytest.raises(ValueError, match="Cannot embed empty batch"): await embedding_service.embed_batch([]) @pytest.mark.asyncio async def test_embed_batch_retry_logic(self, embedding_service, mock_openai_client): """Test retry logic on transient errors.""" texts = ["Test text"] # Mock response for successful retry mock_embedding = MagicMock(spec=Embedding) mock_embedding.embedding = [0.1] * 3072 mock_response = MagicMock(spec=CreateEmbeddingResponse) mock_response.data = [mock_embedding] # Fail twice, then succeed mock_openai_client.embeddings.create = AsyncMock( side_effect=[ Exception("Rate limit exceeded"), Exception("Timeout"), mock_response, ] ) # Call embed_batch embeddings = await embedding_service.embed_batch(texts, max_retries=3) # Should succeed on third attempt assert len(embeddings) == 1 assert len(embeddings[0]) == 3072 assert mock_openai_client.embeddings.create.call_count == 3 @pytest.mark.asyncio async def test_embed_batch_permanent_failure( self, embedding_service, mock_openai_client ): """Test IndexingError raised after max retries.""" texts = ["Test text"] # Always fail mock_openai_client.embeddings.create = AsyncMock( side_effect=Exception("Permanent error") ) # Should raise IndexingError after retries with pytest.raises(IndexingError, match="Failed to embed batch"): await embedding_service.embed_batch(texts, max_retries=3) # Should have tried 3 times assert mock_openai_client.embeddings.create.call_count == 3 @pytest.mark.asyncio async def test_embed_records_success(self, embedding_service, mock_openai_client): """Test embedding ExtractedRecord objects.""" records = [ ExtractedRecord( content={"Item": "A"}, headers=["Item"], _source={ "document_id": str(uuid4()), "filename": "test.xlsx", "sheet": "Sheet1", "table_id": 1, "row": 1, "col_range": "A:A", }, ), ExtractedRecord( content={"Item": "B"}, headers=["Item"], _source={ "document_id": str(uuid4()), "filename": "test.xlsx", "sheet": "Sheet1", "table_id": 1, "row": 2, "col_range": "A:A", }, ), ] # Mock response mock_embedding_1 = MagicMock(spec=Embedding) mock_embedding_1.embedding = [0.1] * 3072 mock_embedding_2 = MagicMock(spec=Embedding) mock_embedding_2.embedding = [0.2] * 3072 mock_response = MagicMock(spec=CreateEmbeddingResponse) mock_response.data = [mock_embedding_1, mock_embedding_2] mock_openai_client.embeddings.create = AsyncMock(return_value=mock_response) # Call embed_records embeddings = await embedding_service.embed_records(records) # Verify results assert len(embeddings) == 2 assert len(embeddings[0]) == 3072 assert len(embeddings[1]) == 3072 # Verify texts were converted properly call_args = mock_openai_client.embeddings.create.call_args input_texts = call_args.kwargs["input"] assert len(input_texts) == 2 assert "Item: A" in input_texts[0] assert "Item: B" in input_texts[1] @pytest.mark.asyncio async def test_embed_records_exceeds_limit(self, embedding_service): """Test error when too many records provided.""" records = [ ExtractedRecord( content={"Item": f"Item{i}"}, headers=["Item"], _source={ "document_id": str(uuid4()), "filename": "test.xlsx", "sheet": "Sheet1", "table_id": 1, "row": i, "col_range": "A:A", }, ) for i in range(101) ] with pytest.raises(ValueError, match="Batch size must be ≤ 100"): await embedding_service.embed_records(records) @pytest.mark.asyncio async def test_embed_batch_timing_metrics( self, embedding_service, mock_openai_client ): """Test that timing metrics are logged.""" texts = ["Test text"] mock_embedding = MagicMock(spec=Embedding) mock_embedding.embedding = [0.1] * 3072 mock_response = MagicMock(spec=CreateEmbeddingResponse) mock_response.data = [mock_embedding] mock_openai_client.embeddings.create = AsyncMock(return_value=mock_response) # Call embed_batch embeddings = await embedding_service.embed_batch(texts) # Verify result assert len(embeddings) == 1 # Timing metrics would be logged (tested via logger assertions in integration tests)