""" Unit tests for StoreVectorsOperator. """ import pytest from unittest.mock import Mock, MagicMock, patch, AsyncMock from uuid import UUID from plugins.operators.store_vectors import StoreVectorsOperator, ExcelRecordProxy class TestStoreVectorsOperator: """Test suite for StoreVectorsOperator.""" @pytest.fixture def mock_context(self): """Create a mock Airflow context.""" return {'ti': Mock()} @pytest.fixture def sample_validate_result(self): """Sample validate result data.""" return { 'bucket': 'documents', 'object_key': 'test/sample.xlsx', 'content_type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', 'document_id': '550e8400-e29b-41d4-a716-446655440000', } @pytest.fixture def sample_embedding_result_excel(self): """Sample embedding result with Excel table chunks.""" return { 'embeddings': [ {'chunk_id': 'doc_t0_r1', 'vector': [0.1] * 3072}, {'chunk_id': 'doc_t0_r2', 'vector': [0.2] * 3072}, ], 'chunks': [ { 'content': {'Name': 'Alice', 'Age': '25'}, '_source': { 'document_id': '550e8400-e29b-41d4-a716-446655440000', 'filename': 'sample.xlsx', 'sheet': 'Sheet1', 'table_id': 0, 'row': 1, 'col_range': 'A:B', }, }, { 'content': {'Name': 'Bob', 'Age': '30'}, '_source': { 'document_id': '550e8400-e29b-41d4-a716-446655440000', 'filename': 'sample.xlsx', 'sheet': 'Sheet1', 'table_id': 0, 'row': 2, 'col_range': 'A:B', }, }, ], 'embedding_count': 2, 'model': 'text-embedding-3-large', } @pytest.fixture def sample_embedding_result_word(self): """Sample embedding result with semantic chunks.""" return { 'embeddings': [ {'chunk_id': '550e8400-e29b-41d4-a716-446655440001', 'vector': [0.1] * 3072}, ], 'chunks': [ { 'text': 'This is the document content.', 'metadata': {'page_number': 1}, 'original_filename': 'document.docx', }, ], 'embedding_count': 1, 'model': 'text-embedding-3-large', } def test_store_vectors_success_excel( self, mock_context, sample_validate_result, sample_embedding_result_excel ): """Test successful vector storage for Excel documents.""" def xcom_pull_side_effect(task_ids): if task_ids == 'generate_embeddings': return sample_embedding_result_excel elif task_ids == 'validate_event': return sample_validate_result return None mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect with patch('plugins.operators.store_vectors.VectorStore') as mock_store_class: # Mock vector store mock_store = MagicMock() mock_store.create_collection = AsyncMock() mock_store.index_with_vectors = AsyncMock(return_value=2) mock_store._disconnect = MagicMock() mock_store.collection_name = 'test_collection' mock_store_class.return_value = mock_store op = StoreVectorsOperator(task_id='test_store_vectors') result = op.execute(mock_context) # Verify result structure assert result['status'] == 'success' assert result['stored_count'] == 2 assert result['document_id'] == '550e8400-e29b-41d4-a716-446655440000' assert result['document_type'] == 'excel' assert result['collection'] == 'test_collection' # Verify VectorStore methods were called mock_store.create_collection.assert_called_once() mock_store.index_with_vectors.assert_called_once() def test_store_vectors_success_word( self, mock_context, sample_validate_result, sample_embedding_result_word ): """Test successful vector storage for Word documents.""" # Update validate result for Word document sample_validate_result['object_key'] = 'test/document.docx' def xcom_pull_side_effect(task_ids): if task_ids == 'generate_embeddings': return sample_embedding_result_word elif task_ids == 'validate_event': return sample_validate_result return None mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect with patch('plugins.operators.store_vectors.VectorStore') as mock_store_class: mock_store = MagicMock() mock_store.create_collection = AsyncMock() mock_store.index_with_vectors = AsyncMock(return_value=1) mock_store._disconnect = MagicMock() mock_store.collection_name = 'test_collection' mock_store_class.return_value = mock_store op = StoreVectorsOperator(task_id='test_store_vectors') result = op.execute(mock_context) assert result['status'] == 'success' assert result['stored_count'] == 1 assert result['document_type'] == 'word' def test_store_vectors_no_embeddings(self, mock_context, sample_validate_result): """Test handling of empty embeddings.""" def xcom_pull_side_effect(task_ids): if task_ids == 'generate_embeddings': return {'embeddings': [], 'chunks': []} elif task_ids == 'validate_event': return sample_validate_result return None mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect op = StoreVectorsOperator(task_id='test_store_vectors') result = op.execute(mock_context) assert result['status'] == 'empty' assert result['stored_count'] == 0 def test_store_vectors_no_embedding_result(self, mock_context, sample_validate_result): """Test handling of missing embedding result.""" def xcom_pull_side_effect(task_ids): if task_ids == 'validate_event': return sample_validate_result return None mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect op = StoreVectorsOperator(task_id='test_store_vectors') result = op.execute(mock_context) assert result['status'] == 'empty' assert result['stored_count'] == 0 def test_store_vectors_missing_validate_result_raises(self, mock_context): """Test that missing validate result raises ValueError.""" mock_context['ti'].xcom_pull.return_value = None op = StoreVectorsOperator(task_id='test_store_vectors') with pytest.raises(ValueError, match='No validate result'): op.execute(mock_context) def test_custom_task_ids(self): """Test using custom task IDs.""" op = StoreVectorsOperator( task_id='test_store', embeddings_task_id='custom_embeddings', validate_task_id='custom_validate', ) assert op.embeddings_task_id == 'custom_embeddings' assert op.validate_task_id == 'custom_validate' def test_default_constants(self): """Test default constant values.""" assert StoreVectorsOperator.DEFAULT_EMBEDDINGS_TASK_ID == 'generate_embeddings' assert StoreVectorsOperator.DEFAULT_VALIDATE_TASK_ID == 'validate_event' def test_empty_result(self): """Test _empty_result helper method.""" op = StoreVectorsOperator(task_id='test') doc_id = UUID('550e8400-e29b-41d4-a716-446655440000') result = op._empty_result(doc_id) assert result['stored_count'] == 0 assert result['document_id'] == '550e8400-e29b-41d4-a716-446655440000' assert result['status'] == 'empty' assert result['collection'] == '' assert result['document_type'] == 'unknown' def test_detect_document_type_excel(self): """Test document type detection for Excel chunks.""" op = StoreVectorsOperator(task_id='test') chunks = [ { 'content': {'Name': 'Alice'}, '_source': {'sheet': 'Sheet1', 'row': 1}, } ] doc_type = op._detect_document_type(chunks) assert doc_type == 'excel' def test_detect_document_type_word(self): """Test document type detection for Word chunks.""" op = StoreVectorsOperator(task_id='test') chunks = [ { 'text': 'Document content', 'metadata': {'page_number': 1}, } ] doc_type = op._detect_document_type(chunks) assert doc_type == 'word' def test_detect_document_type_content_sheet(self): """Test document type detection for 'content' sheet (semantic).""" op = StoreVectorsOperator(task_id='test') chunks = [ { 'content': {'text': 'Some text'}, '_source': {'sheet': 'content'}, } ] doc_type = op._detect_document_type(chunks) assert doc_type == 'word' def test_detect_document_type_empty(self): """Test document type detection for empty chunks.""" op = StoreVectorsOperator(task_id='test') doc_type = op._detect_document_type([]) assert doc_type == 'word' # Default def test_generate_record_ids_with_valid_uuids(self): """Test record ID generation with valid UUID chunk IDs.""" op = StoreVectorsOperator(task_id='test') embeddings = [ {'chunk_id': '550e8400-e29b-41d4-a716-446655440001', 'vector': []}, {'chunk_id': '550e8400-e29b-41d4-a716-446655440002', 'vector': []}, ] ids = op._generate_record_ids(embeddings, 2) assert len(ids) == 2 assert ids[0] == UUID('550e8400-e29b-41d4-a716-446655440001') assert ids[1] == UUID('550e8400-e29b-41d4-a716-446655440002') def test_generate_record_ids_with_invalid_uuids(self): """Test record ID generation with non-UUID chunk IDs.""" op = StoreVectorsOperator(task_id='test') embeddings = [ {'chunk_id': 'doc_t0_r1', 'vector': []}, # Not a valid UUID ] ids = op._generate_record_ids(embeddings, 1) assert len(ids) == 1 # Should generate a new UUID, not raise an error assert isinstance(ids[0], UUID) def test_prepare_records_for_milvus_excel(self): """Test record preparation for Excel documents.""" op = StoreVectorsOperator(task_id='test') chunks = [ { 'content': {'Name': 'Alice'}, '_source': {'sheet': 'Sheet1', 'row': 1, 'filename': 'test.xlsx'}, } ] records = op._prepare_records_for_milvus(chunks, 'excel') assert len(records) == 1 assert isinstance(records[0], ExcelRecordProxy) assert records[0].content == {'Name': 'Alice'} def test_prepare_records_for_milvus_word(self): """Test record preparation for Word documents.""" op = StoreVectorsOperator(task_id='test') chunks = [ { 'text': 'Document content', 'metadata': {'page_number': 1}, } ] records = op._prepare_records_for_milvus(chunks, 'word') assert len(records) == 1 assert isinstance(records[0], dict) assert records[0]['text'] == 'Document content' def test_prepare_records_for_milvus_word_from_content(self): """Test record preparation extracts text from content dict.""" op = StoreVectorsOperator(task_id='test') chunks = [ { 'content': {'text': 'Content from dict'}, '_source': {'sheet': 'content'}, } ] records = op._prepare_records_for_milvus(chunks, 'word') assert len(records) == 1 assert records[0]['text'] == 'Content from dict' def test_store_vectors_error_handling( self, mock_context, sample_validate_result, sample_embedding_result_excel ): """Test that VectorStore errors are properly propagated.""" def xcom_pull_side_effect(task_ids): if task_ids == 'generate_embeddings': return sample_embedding_result_excel elif task_ids == 'validate_event': return sample_validate_result return None mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect with patch('plugins.operators.store_vectors.VectorStore') as mock_store_class: mock_store = MagicMock() mock_store.create_collection = AsyncMock() mock_store.index_with_vectors = AsyncMock( side_effect=Exception('Milvus connection failed') ) mock_store._disconnect = MagicMock() mock_store.collection_name = 'test_collection' mock_store_class.return_value = mock_store op = StoreVectorsOperator(task_id='test_store_vectors') with pytest.raises(Exception, match='Milvus connection failed'): op.execute(mock_context) # Verify disconnect was called even on error mock_store._disconnect.assert_called_once() def test_store_vectors_disconnects_on_success( self, mock_context, sample_validate_result, sample_embedding_result_excel ): """Test that VectorStore disconnects after successful operation.""" def xcom_pull_side_effect(task_ids): if task_ids == 'generate_embeddings': return sample_embedding_result_excel elif task_ids == 'validate_event': return sample_validate_result return None mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect with patch('plugins.operators.store_vectors.VectorStore') as mock_store_class: mock_store = MagicMock() mock_store.create_collection = AsyncMock() mock_store.index_with_vectors = AsyncMock(return_value=2) mock_store._disconnect = MagicMock() mock_store.collection_name = 'test_collection' mock_store_class.return_value = mock_store op = StoreVectorsOperator(task_id='test_store_vectors') result = op.execute(mock_context) assert result['status'] == 'success' # Verify disconnect was called mock_store._disconnect.assert_called_once() class TestExcelRecordProxy: """Test suite for ExcelRecordProxy helper class.""" def test_proxy_from_full_chunk(self): """Test proxy creation from complete chunk.""" chunk = { 'content': {'Name': 'Alice', 'Age': '25'}, '_source': { 'sheet': 'Sheet1', 'row': 1, 'filename': 'test.xlsx', }, 'resolved_content': {'Name': 'Alice Smith'}, } proxy = ExcelRecordProxy(chunk) assert proxy.content == {'Name': 'Alice', 'Age': '25'} assert proxy.headers == ['Name', 'Age'] assert proxy._source['sheet'] == 'Sheet1' assert proxy.resolved_content == {'Name': 'Alice Smith'} def test_proxy_from_minimal_chunk(self): """Test proxy creation from minimal chunk.""" chunk = { 'content': {'text': 'Some text'}, } proxy = ExcelRecordProxy(chunk) assert proxy.content == {'text': 'Some text'} assert proxy.headers == ['text'] # Should have default _source values assert proxy._source.get('sheet') == 'Unknown' assert proxy._source.get('row') == 0 assert proxy._source.get('filename') == 'unknown' assert proxy.resolved_content is None def test_proxy_headers_from_non_dict_content(self): """Test proxy handles non-dict content.""" chunk = { 'content': 'Just a string', '_source': {'sheet': 'Sheet1'}, } proxy = ExcelRecordProxy(chunk) assert proxy.content == 'Just a string' assert proxy.headers == []