From f7437cf72d90e9d4ce7bf71ddc597389ac05eff7 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 14 Jul 2025 15:55:12 +0100 Subject: [PATCH] Testing embeedings --- tests/unit/test_embeddings/__init__.py | 10 + tests/unit/test_embeddings/conftest.py | 114 +++++ .../test_embeddings/test_embedding_logic.py | 278 ++++++++++++ .../test_embeddings/test_embedding_utils.py | 349 +++++++++++++++ .../test_embeddings/test_fastembed_core.py | 213 +++++++++ .../test_fastembed_processor.py | 403 ++++++++++++++++++ .../unit/test_embeddings/test_ollama_core.py | 222 ++++++++++ .../test_embeddings/test_ollama_processor.py | 334 +++++++++++++++ 8 files changed, 1923 insertions(+) create mode 100644 tests/unit/test_embeddings/__init__.py create mode 100644 tests/unit/test_embeddings/conftest.py create mode 100644 tests/unit/test_embeddings/test_embedding_logic.py create mode 100644 tests/unit/test_embeddings/test_embedding_utils.py create mode 100644 tests/unit/test_embeddings/test_fastembed_core.py create mode 100644 tests/unit/test_embeddings/test_fastembed_processor.py create mode 100644 tests/unit/test_embeddings/test_ollama_core.py create mode 100644 tests/unit/test_embeddings/test_ollama_processor.py diff --git a/tests/unit/test_embeddings/__init__.py b/tests/unit/test_embeddings/__init__.py new file mode 100644 index 00000000..9320e90f --- /dev/null +++ b/tests/unit/test_embeddings/__init__.py @@ -0,0 +1,10 @@ +""" +Unit tests for embeddings services + +Testing Strategy: +- Mock external embedding libraries (FastEmbed, Ollama client) +- Test core business logic for text embedding generation +- Test error handling and edge cases +- Test vector dimension consistency +- Test batch processing logic +""" \ No newline at end of file diff --git a/tests/unit/test_embeddings/conftest.py b/tests/unit/test_embeddings/conftest.py new file mode 100644 index 00000000..ac1346eb --- /dev/null +++ b/tests/unit/test_embeddings/conftest.py @@ -0,0 +1,114 @@ +""" +Shared fixtures for embeddings unit tests +""" + +import pytest +import numpy as np +from unittest.mock import Mock, AsyncMock, MagicMock +from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error + + +@pytest.fixture +def sample_text(): + """Sample text for embedding tests""" + return "This is a sample text for embedding generation." + + +@pytest.fixture +def sample_embedding_vector(): + """Sample embedding vector for mocking""" + return [0.1, 0.2, -0.3, 0.4, -0.5, 0.6, 0.7, -0.8, 0.9, -1.0] + + +@pytest.fixture +def sample_batch_embeddings(): + """Sample batch of embedding vectors""" + return [ + [0.1, 0.2, -0.3, 0.4, -0.5], + [0.6, 0.7, -0.8, 0.9, -1.0], + [-0.1, -0.2, 0.3, -0.4, 0.5] + ] + + +@pytest.fixture +def sample_embeddings_request(): + """Sample EmbeddingsRequest for testing""" + return EmbeddingsRequest( + text="Test text for embedding" + ) + + +@pytest.fixture +def sample_embeddings_response(sample_embedding_vector): + """Sample successful EmbeddingsResponse""" + return EmbeddingsResponse( + error=None, + vectors=sample_embedding_vector + ) + + +@pytest.fixture +def sample_error_response(): + """Sample error EmbeddingsResponse""" + return EmbeddingsResponse( + error=Error(type="embedding-error", message="Model not found"), + vectors=None + ) + + +@pytest.fixture +def mock_message(): + """Mock Pulsar message for testing""" + message = Mock() + message.properties.return_value = {"id": "test-message-123"} + return message + + +@pytest.fixture +def mock_flow(): + """Mock flow for producer/consumer testing""" + flow = Mock() + flow.return_value.send = AsyncMock() + flow.producer = {"response": Mock()} + flow.producer["response"].send = AsyncMock() + return flow + + +@pytest.fixture +def mock_consumer(): + """Mock Pulsar consumer""" + return AsyncMock() + + +@pytest.fixture +def mock_producer(): + """Mock Pulsar producer""" + return AsyncMock() + + +@pytest.fixture +def mock_fastembed_embedding(): + """Mock FastEmbed TextEmbedding""" + mock = Mock() + mock.embed.return_value = [np.array([0.1, 0.2, -0.3, 0.4, -0.5])] + return mock + + +@pytest.fixture +def mock_ollama_client(): + """Mock Ollama client""" + mock = Mock() + mock.embed.return_value = Mock( + embeddings=[0.1, 0.2, -0.3, 0.4, -0.5] + ) + return mock + + +@pytest.fixture +def embedding_test_params(): + """Common parameters for embedding processor testing""" + return { + "model": "test-model", + "concurrency": 1, + "id": "test-embeddings" + } \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_embedding_logic.py b/tests/unit/test_embeddings/test_embedding_logic.py new file mode 100644 index 00000000..055cb2d1 --- /dev/null +++ b/tests/unit/test_embeddings/test_embedding_logic.py @@ -0,0 +1,278 @@ +""" +Unit tests for embedding business logic + +Tests the core embedding functionality without external dependencies, +focusing on data processing, validation, and business rules. +""" + +import pytest +import numpy as np +from unittest.mock import Mock, patch + + +class TestEmbeddingBusinessLogic: + """Test embedding business logic and data processing""" + + def test_embedding_vector_validation(self): + """Test validation of embedding vectors""" + # Arrange + valid_vectors = [ + [0.1, 0.2, 0.3], + [-0.5, 0.0, 0.8], + [], # Empty vector + [1.0] * 1536 # Large vector + ] + + invalid_vectors = [ + None, + "not a vector", + [1, 2, "string"], + [[1, 2], [3, 4]] # Nested + ] + + # Act & Assert + def is_valid_vector(vec): + if not isinstance(vec, list): + return False + return all(isinstance(x, (int, float)) for x in vec) + + for vec in valid_vectors: + assert is_valid_vector(vec), f"Should be valid: {vec}" + + for vec in invalid_vectors: + assert not is_valid_vector(vec), f"Should be invalid: {vec}" + + def test_dimension_consistency_check(self): + """Test dimension consistency validation""" + # Arrange + same_dimension_vectors = [ + [0.1, 0.2, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9, 1.0], + [-0.1, -0.2, -0.3, -0.4, -0.5] + ] + + mixed_dimension_vectors = [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6, 0.7], + [0.8, 0.9] + ] + + # Act + def check_dimension_consistency(vectors): + if not vectors: + return True + expected_dim = len(vectors[0]) + return all(len(vec) == expected_dim for vec in vectors) + + # Assert + assert check_dimension_consistency(same_dimension_vectors) + assert not check_dimension_consistency(mixed_dimension_vectors) + + def test_text_preprocessing_logic(self): + """Test text preprocessing for embeddings""" + # Arrange + test_cases = [ + ("Simple text", "Simple text"), + ("", ""), + ("Text with\nnewlines", "Text with\nnewlines"), + ("Unicode: 世界 🌍", "Unicode: 世界 🌍"), + (" Whitespace ", " Whitespace ") + ] + + # Act & Assert + for input_text, expected in test_cases: + # Simple preprocessing (identity in this case) + processed = str(input_text) if input_text is not None else "" + assert processed == expected + + def test_batch_processing_logic(self): + """Test batch processing logic for multiple texts""" + # Arrange + texts = ["Text 1", "Text 2", "Text 3"] + + def mock_embed_single(text): + # Simulate embedding generation based on text length + return [len(text) / 10.0] * 5 + + # Act + results = [] + for text in texts: + embedding = mock_embed_single(text) + results.append((text, embedding)) + + # Assert + assert len(results) == len(texts) + for i, (original_text, embedding) in enumerate(results): + assert original_text == texts[i] + assert len(embedding) == 5 + expected_value = len(texts[i]) / 10.0 + assert all(abs(val - expected_value) < 0.001 for val in embedding) + + def test_numpy_array_conversion_logic(self): + """Test numpy array to list conversion""" + # Arrange + test_arrays = [ + np.array([1, 2, 3], dtype=np.int32), + np.array([1.0, 2.0, 3.0], dtype=np.float64), + np.array([0.1, 0.2, 0.3], dtype=np.float32) + ] + + # Act + converted = [] + for arr in test_arrays: + result = arr.tolist() + converted.append(result) + + # Assert + assert converted[0] == [1, 2, 3] + assert converted[1] == [1.0, 2.0, 3.0] + # Float32 might have precision differences, so check approximately + assert len(converted[2]) == 3 + assert all(isinstance(x, float) for x in converted[2]) + + def test_error_response_generation(self): + """Test error response generation logic""" + # Arrange + error_scenarios = [ + ("model_not_found", "Model 'xyz' not found"), + ("connection_error", "Failed to connect to service"), + ("rate_limit", "Rate limit exceeded"), + ("invalid_input", "Invalid input format") + ] + + # Act & Assert + for error_type, error_message in error_scenarios: + error_response = { + "error": { + "type": error_type, + "message": error_message + }, + "vectors": None + } + + assert error_response["error"]["type"] == error_type + assert error_response["error"]["message"] == error_message + assert error_response["vectors"] is None + + def test_success_response_generation(self): + """Test success response generation logic""" + # Arrange + test_vectors = [0.1, 0.2, 0.3, 0.4, 0.5] + + # Act + success_response = { + "error": None, + "vectors": test_vectors + } + + # Assert + assert success_response["error"] is None + assert success_response["vectors"] == test_vectors + assert len(success_response["vectors"]) == 5 + + def test_model_parameter_handling(self): + """Test model parameter validation and handling""" + # Arrange + valid_models = { + "ollama": ["mxbai-embed-large", "nomic-embed-text"], + "fastembed": ["sentence-transformers/all-MiniLM-L6-v2", "BAAI/bge-small-en-v1.5"] + } + + # Act & Assert + for provider, models in valid_models.items(): + for model in models: + assert isinstance(model, str) + assert len(model) > 0 + if provider == "fastembed": + assert "/" in model or "-" in model + + def test_concurrent_processing_simulation(self): + """Test concurrent processing simulation""" + # Arrange + import asyncio + + async def mock_async_embed(text, delay=0.001): + await asyncio.sleep(delay) + return [ord(text[0]) / 255.0] if text else [0.0] + + # Act + async def run_concurrent(): + texts = ["A", "B", "C", "D", "E"] + tasks = [mock_async_embed(text) for text in texts] + results = await asyncio.gather(*tasks) + return list(zip(texts, results)) + + # Run test + results = asyncio.run(run_concurrent()) + + # Assert + assert len(results) == 5 + for i, (text, embedding) in enumerate(results): + expected_char = chr(ord('A') + i) + assert text == expected_char + expected_value = ord(expected_char) / 255.0 + assert abs(embedding[0] - expected_value) < 0.001 + + def test_empty_and_edge_cases(self): + """Test empty inputs and edge cases""" + # Arrange + edge_cases = [ + ("", "empty string"), + (" ", "single space"), + ("a", "single character"), + ("A" * 10000, "very long string"), + ("\\n\\t\\r", "special characters"), + ("混合English中文", "mixed languages") + ] + + # Act & Assert + for text, description in edge_cases: + # Basic validation that text can be processed + assert isinstance(text, str), f"Failed for {description}" + assert len(text) >= 0, f"Failed for {description}" + + # Simulate embedding generation would work + mock_embedding = [len(text) % 10] * 3 + assert len(mock_embedding) == 3, f"Failed for {description}" + + def test_vector_normalization_logic(self): + """Test vector normalization calculations""" + # Arrange + test_vectors = [ + [3.0, 4.0], # Should normalize to [0.6, 0.8] + [1.0, 0.0], # Should normalize to [1.0, 0.0] + [0.0, 0.0], # Zero vector edge case + ] + + # Act & Assert + for vector in test_vectors: + magnitude = sum(x**2 for x in vector) ** 0.5 + + if magnitude > 0: + normalized = [x / magnitude for x in vector] + # Check unit length (approximately) + norm_magnitude = sum(x**2 for x in normalized) ** 0.5 + assert abs(norm_magnitude - 1.0) < 0.0001 + else: + # Zero vector case + assert all(x == 0 for x in vector) + + def test_cosine_similarity_calculation(self): + """Test cosine similarity computation""" + # Arrange + vector_pairs = [ + ([1, 0], [0, 1], 0.0), # Orthogonal + ([1, 0], [1, 0], 1.0), # Identical + ([1, 1], [-1, -1], -1.0), # Opposite + ] + + # Act & Assert + def cosine_similarity(v1, v2): + dot = sum(a * b for a, b in zip(v1, v2)) + mag1 = sum(x**2 for x in v1) ** 0.5 + mag2 = sum(x**2 for x in v2) ** 0.5 + return dot / (mag1 * mag2) if mag1 * mag2 > 0 else 0 + + for v1, v2, expected in vector_pairs: + similarity = cosine_similarity(v1, v2) + assert abs(similarity - expected) < 0.0001 \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_embedding_utils.py b/tests/unit/test_embeddings/test_embedding_utils.py new file mode 100644 index 00000000..425da28b --- /dev/null +++ b/tests/unit/test_embeddings/test_embedding_utils.py @@ -0,0 +1,349 @@ +""" +Unit tests for embedding utilities and common functionality + +Tests dimension consistency, batch processing, error handling patterns, +and other utilities common across embedding services. +""" + +import pytest +from unittest.mock import patch, Mock, AsyncMock +import numpy as np + +from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error +from trustgraph.exceptions import TooManyRequests + + +class MockEmbeddingProcessor: + """Simple mock embedding processor for testing functionality""" + + def __init__(self, embedding_function=None, **params): + # Store embedding function for mocking + self.embedding_function = embedding_function + self.model = params.get('model', 'test-model') + + async def on_embeddings(self, text): + if self.embedding_function: + return self.embedding_function(text) + return [0.1, 0.2, 0.3, 0.4, 0.5] # Default test embedding + + +class TestEmbeddingDimensionConsistency: + """Test cases for embedding dimension consistency""" + + async def test_consistent_dimensions_single_processor(self): + """Test that a single processor returns consistent dimensions""" + # Arrange + dimension = 128 + def mock_embedding(text): + return [0.1] * dimension + + processor = MockEmbeddingProcessor(embedding_function=mock_embedding) + + # Act + results = [] + test_texts = ["Text 1", "Text 2", "Text 3", "Text 4", "Text 5"] + + for text in test_texts: + result = await processor.on_embeddings(text) + results.append(result) + + # Assert + for result in results: + assert len(result) == dimension, f"Expected dimension {dimension}, got {len(result)}" + + # All results should have same dimensions + first_dim = len(results[0]) + for i, result in enumerate(results[1:], 1): + assert len(result) == first_dim, f"Dimension mismatch at index {i}" + + async def test_dimension_consistency_across_text_lengths(self): + """Test dimension consistency across varying text lengths""" + # Arrange + dimension = 384 + def mock_embedding(text): + # Dimension should not depend on text length + return [0.1] * dimension + + processor = MockEmbeddingProcessor(embedding_function=mock_embedding) + + # Act - Test various text lengths + test_texts = [ + "", # Empty text + "Hi", # Very short + "This is a medium length sentence for testing.", # Medium + "This is a very long text that should still produce embeddings of consistent dimension regardless of the input text length and content." * 10 # Very long + ] + + results = [] + for text in test_texts: + result = await processor.on_embeddings(text) + results.append(result) + + # Assert + for i, result in enumerate(results): + assert len(result) == dimension, f"Text length {len(test_texts[i])} produced wrong dimension" + + async def test_dimension_validation_different_models(self): + """Test dimension validation for different model configurations""" + # Arrange + models_and_dims = [ + ("small-model", 128), + ("medium-model", 384), + ("large-model", 1536) + ] + + for model_name, expected_dim in models_and_dims: + def mock_embedding(text): + return [0.1] * expected_dim + + processor = MockEmbeddingProcessor( + model=model_name, + embedding_function=mock_embedding + ) + + # Act + result = await processor.on_embeddings("Test text") + + # Assert + assert len(result) == expected_dim, f"Model {model_name} produced wrong dimension" + + +class TestEmbeddingBatchProcessing: + """Test cases for batch processing logic""" + + async def test_sequential_processing_maintains_order(self): + """Test that sequential processing maintains text order""" + # Arrange + def mock_embedding(text): + # Return embedding that encodes the text for verification + return [ord(text[0]) / 255.0] if text else [0.0] # Normalize to [0,1] + + processor = MockEmbeddingProcessor(embedding_function=mock_embedding) + + # Act + test_texts = ["A", "B", "C", "D", "E"] + results = [] + + for text in test_texts: + result = await processor.on_embeddings(text) + results.append((text, result)) + + # Assert + for i, (original_text, embedding) in enumerate(results): + assert original_text == test_texts[i] + expected_value = ord(test_texts[i][0]) / 255.0 + assert abs(embedding[0] - expected_value) < 0.001 + + async def test_batch_processing_throughput(self): + """Test batch processing capabilities""" + # Arrange + call_count = 0 + def mock_embedding(text): + nonlocal call_count + call_count += 1 + return [0.1, 0.2, 0.3] + + processor = MockEmbeddingProcessor(embedding_function=mock_embedding) + + # Act - Process multiple texts + batch_size = 10 + test_texts = [f"Text {i}" for i in range(batch_size)] + + results = [] + for text in test_texts: + result = await processor.on_embeddings(text) + results.append(result) + + # Assert + assert call_count == batch_size + assert len(results) == batch_size + for result in results: + assert result == [0.1, 0.2, 0.3] + + async def test_concurrent_processing_simulation(self): + """Test concurrent processing behavior simulation""" + # Arrange + import asyncio + + processing_times = [] + def mock_embedding(text): + import time + processing_times.append(time.time()) + return [len(text) / 100.0] # Encoding text length + + processor = MockEmbeddingProcessor(embedding_function=mock_embedding) + + # Act - Simulate concurrent processing + test_texts = [f"Text {i}" for i in range(5)] + + tasks = [processor.on_embeddings(text) for text in test_texts] + results = await asyncio.gather(*tasks) + + # Assert + assert len(results) == 5 + assert len(processing_times) == 5 + + # Results should correspond to text lengths + for i, result in enumerate(results): + expected_value = len(test_texts[i]) / 100.0 + assert abs(result[0] - expected_value) < 0.001 + + +class TestEmbeddingErrorHandling: + """Test cases for error handling in embedding services""" + + async def test_embedding_function_error_handling(self): + """Test error handling in embedding function""" + # Arrange + def failing_embedding(text): + raise Exception("Embedding model failed") + + processor = MockEmbeddingProcessor(embedding_function=failing_embedding) + + # Act & Assert + with pytest.raises(Exception, match="Embedding model failed"): + await processor.on_embeddings("Test text") + + async def test_rate_limit_exception_propagation(self): + """Test that rate limit exceptions are properly propagated""" + # Arrange + def rate_limited_embedding(text): + raise TooManyRequests("Rate limit exceeded") + + processor = MockEmbeddingProcessor(embedding_function=rate_limited_embedding) + + # Act & Assert + with pytest.raises(TooManyRequests, match="Rate limit exceeded"): + await processor.on_embeddings("Test text") + + async def test_none_result_handling(self): + """Test handling when embedding function returns None""" + # Arrange + def none_embedding(text): + return None + + processor = MockEmbeddingProcessor(embedding_function=none_embedding) + + # Act + result = await processor.on_embeddings("Test text") + + # Assert + assert result is None + + async def test_invalid_embedding_format_handling(self): + """Test handling of invalid embedding formats""" + # Arrange + def invalid_embedding(text): + return "not a list" # Invalid format + + processor = MockEmbeddingProcessor(embedding_function=invalid_embedding) + + # Act + result = await processor.on_embeddings("Test text") + + # Assert + assert result == "not a list" # Returns what the function provides + + +class TestEmbeddingUtilities: + """Test cases for embedding utility functions and helpers""" + + def test_vector_normalization_simulation(self): + """Test vector normalization logic simulation""" + # Arrange + test_vectors = [ + [1.0, 2.0, 3.0], + [0.5, -0.5, 1.0], + [10.0, 20.0, 30.0] + ] + + # Act - Simulate L2 normalization + normalized_vectors = [] + for vector in test_vectors: + magnitude = sum(x**2 for x in vector) ** 0.5 + if magnitude > 0: + normalized = [x / magnitude for x in vector] + else: + normalized = vector + normalized_vectors.append(normalized) + + # Assert + for normalized in normalized_vectors: + magnitude = sum(x**2 for x in normalized) ** 0.5 + assert abs(magnitude - 1.0) < 0.0001, "Vector should be unit length" + + def test_cosine_similarity_calculation(self): + """Test cosine similarity calculation between embeddings""" + # Arrange + vector1 = [1.0, 0.0, 0.0] + vector2 = [0.0, 1.0, 0.0] + vector3 = [1.0, 0.0, 0.0] # Same as vector1 + + # Act - Calculate cosine similarities + def cosine_similarity(v1, v2): + dot_product = sum(a * b for a, b in zip(v1, v2)) + mag1 = sum(x**2 for x in v1) ** 0.5 + mag2 = sum(x**2 for x in v2) ** 0.5 + return dot_product / (mag1 * mag2) if mag1 * mag2 > 0 else 0 + + sim_12 = cosine_similarity(vector1, vector2) + sim_13 = cosine_similarity(vector1, vector3) + + # Assert + assert abs(sim_12 - 0.0) < 0.0001, "Orthogonal vectors should have 0 similarity" + assert abs(sim_13 - 1.0) < 0.0001, "Identical vectors should have 1.0 similarity" + + def test_embedding_validation_helpers(self): + """Test embedding validation helper functions""" + # Arrange + valid_embeddings = [ + [0.1, 0.2, 0.3], + [1.0, -1.0, 0.0], + [] # Empty embedding + ] + + invalid_embeddings = [ + None, + "not a list", + [1, 2, "three"], # Mixed types + [[1, 2], [3, 4]] # Nested lists + ] + + # Act & Assert + def is_valid_embedding(embedding): + if not isinstance(embedding, list): + return False + return all(isinstance(x, (int, float)) for x in embedding) + + for embedding in valid_embeddings: + assert is_valid_embedding(embedding), f"Should be valid: {embedding}" + + for embedding in invalid_embeddings: + assert not is_valid_embedding(embedding), f"Should be invalid: {embedding}" + + async def test_embedding_metadata_handling(self): + """Test handling of embedding metadata and properties""" + # Arrange + def metadata_embedding(text): + return { + "vectors": [0.1, 0.2, 0.3], + "model": "test-model", + "dimension": 3, + "text_length": len(text) + } + + # Mock processor that returns metadata + class MetadataProcessor(MockEmbeddingProcessor): + async def on_embeddings(self, text): + result = metadata_embedding(text) + return result["vectors"] # Return only vectors for compatibility + + processor = MetadataProcessor() + + # Act + result = await processor.on_embeddings("Test text with metadata") + + # Assert + assert isinstance(result, list) + assert len(result) == 3 + assert result == [0.1, 0.2, 0.3] \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_fastembed_core.py b/tests/unit/test_embeddings/test_fastembed_core.py new file mode 100644 index 00000000..b10042dc --- /dev/null +++ b/tests/unit/test_embeddings/test_fastembed_core.py @@ -0,0 +1,213 @@ +""" +Unit tests for FastEmbed core embedding functionality + +Tests the core business logic without full processor initialization, +focusing on the embedding generation methods. +""" + +import pytest +from unittest.mock import patch, Mock +import numpy as np + + +class TestFastEmbedCore: + """Test core FastEmbed embedding functionality""" + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + def test_fastembed_initialization(self, mock_text_embedding): + """Test FastEmbed TextEmbedding initialization""" + # Arrange + mock_embedding = Mock() + mock_text_embedding.return_value = mock_embedding + + # Act + from trustgraph.embeddings.fastembed.processor import TextEmbedding + embedding_model = TextEmbedding(model_name="test-model") + + # Assert + mock_text_embedding.assert_called_once_with(model_name="test-model") + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + def test_fastembed_embed_method(self, mock_text_embedding): + """Test FastEmbed embed method functionality""" + # Arrange + mock_embedding = Mock() + mock_embedding.embed.return_value = [ + np.array([0.1, 0.2, -0.3, 0.4, -0.5]) + ] + mock_text_embedding.return_value = mock_embedding + + # Act + from trustgraph.embeddings.fastembed.processor import TextEmbedding + embedding_model = TextEmbedding(model_name="test-model") + result = embedding_model.embed(["Test text"]) + + # Assert + mock_embedding.embed.assert_called_once_with(["Test text"]) + assert len(result) == 1 + np.testing.assert_array_equal(result[0], np.array([0.1, 0.2, -0.3, 0.4, -0.5])) + + def test_numpy_to_list_conversion(self): + """Test numpy array to list conversion logic""" + # Arrange + test_arrays = [ + np.array([0.1, 0.2, 0.3], dtype=np.float32), + np.array([-0.4, 0.5, -0.6], dtype=np.float64), + np.array([1, 2, 3], dtype=np.int32) + ] + + # Act + converted = [arr.tolist() for arr in test_arrays] + + # Assert + assert all(isinstance(vec, list) for vec in converted) + assert len(converted) == 3 + + # Check float32 conversion (with tolerance for precision) + assert len(converted[0]) == 3 + assert abs(converted[0][0] - 0.1) < 0.001 + assert abs(converted[0][1] - 0.2) < 0.001 + assert abs(converted[0][2] - 0.3) < 0.001 + + # Check float64 conversion (more precise) + assert converted[1] == [-0.4, 0.5, -0.6] + + # Check integer conversion (exact) + assert converted[2] == [1, 2, 3] + + assert all(isinstance(val, (int, float)) for vec in converted for val in vec) + + def test_embedding_dimension_consistency(self): + """Test that embeddings maintain consistent dimensions""" + # Arrange + test_vectors = [ + np.array([0.1, 0.2, 0.3, 0.4, 0.5]), + np.array([0.6, 0.7, 0.8, 0.9, 1.0]), + np.array([-0.1, -0.2, -0.3, -0.4, -0.5]) + ] + + # Act + converted = [vec.tolist() for vec in test_vectors] + + # Assert + dimensions = [len(vec) for vec in converted] + assert all(dim == dimensions[0] for dim in dimensions), "All vectors should have same dimension" + assert dimensions[0] == 5 + + def test_embedding_batch_processing_logic(self): + """Test batch processing logic for multiple texts""" + # Arrange + test_texts = ["Text 1", "Text 2", "Text 3"] + mock_embeddings = [ + np.array([0.1, 0.2, 0.3]), + np.array([0.4, 0.5, 0.6]), + np.array([0.7, 0.8, 0.9]) + ] + + # Act - Simulate the conversion logic used in FastEmbed processor + result = [vec.tolist() for vec in mock_embeddings] + + # Assert + assert len(result) == len(test_texts) + assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]] + + def test_empty_text_handling(self): + """Test handling of empty text input""" + # Arrange + empty_embedding = np.array([]) + + # Act + result = empty_embedding.tolist() + + # Assert + assert result == [] + assert isinstance(result, list) + + def test_large_embedding_vectors(self): + """Test handling of large embedding vectors""" + # Arrange + large_dimension = 1536 + large_vector = np.random.rand(large_dimension) + + # Act + result = large_vector.tolist() + + # Assert + assert len(result) == large_dimension + assert isinstance(result, list) + assert all(isinstance(val, float) for val in result) + + def test_unicode_text_compatibility(self): + """Test that Unicode text is properly handled""" + # Arrange + unicode_texts = [ + "Hello 世界", + "Café naïve", + "🚀 Rocket", + "Ελληνικά", + "русский" + ] + + # Act - Simulate text processing + processed_texts = [text.encode('utf-8').decode('utf-8') for text in unicode_texts] + + # Assert + assert processed_texts == unicode_texts + assert all(isinstance(text, str) for text in processed_texts) + + def test_model_parameter_validation(self): + """Test model parameter validation""" + # Arrange + valid_models = [ + "sentence-transformers/all-MiniLM-L6-v2", + "sentence-transformers/paraphrase-MiniLM-L6-v2", + "BAAI/bge-small-en-v1.5" + ] + + # Act & Assert - All should be valid string parameters + for model in valid_models: + assert isinstance(model, str) + assert len(model) > 0 + assert "/" in model or "-" in model # Common model naming patterns + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + def test_error_handling_simulation(self, mock_text_embedding): + """Test error handling in embedding generation""" + # Arrange + mock_embedding = Mock() + mock_embedding.embed.side_effect = Exception("Model loading failed") + mock_text_embedding.return_value = mock_embedding + + # Act & Assert + from trustgraph.embeddings.fastembed.processor import TextEmbedding + embedding_model = TextEmbedding(model_name="test-model") + + with pytest.raises(Exception, match="Model loading failed"): + embedding_model.embed(["Test text"]) + + def test_default_model_constant(self): + """Test that default model constant is properly defined""" + # Act + from trustgraph.embeddings.fastembed.processor import default_model + + # Assert + assert default_model == "sentence-transformers/all-MiniLM-L6-v2" + assert isinstance(default_model, str) + assert len(default_model) > 0 + + def test_processor_add_args_functionality(self): + """Test add_args static method functionality""" + # Arrange + from trustgraph.embeddings.fastembed.processor import Processor + mock_parser = Mock() + + # Act + Processor.add_args(mock_parser) + + # Assert + # Verify model argument was added + model_calls = [call for call in mock_parser.add_argument.call_args_list + if call[0][0] in ['-m', '--model']] + assert len(model_calls) == 1 + model_call = model_calls[0] + assert model_call[1]['default'] == "sentence-transformers/all-MiniLM-L6-v2" \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_fastembed_processor.py b/tests/unit/test_embeddings/test_fastembed_processor.py new file mode 100644 index 00000000..c297e685 --- /dev/null +++ b/tests/unit/test_embeddings/test_fastembed_processor.py @@ -0,0 +1,403 @@ +""" +Unit tests for FastEmbed embeddings processor + +Tests the core business logic of FastEmbed embedding generation without +relying on external FastEmbed library infrastructure. +""" + +import pytest +from unittest.mock import patch, Mock, AsyncMock, MagicMock +import numpy as np + +from trustgraph.embeddings.fastembed.processor import Processor +from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error +from trustgraph.exceptions import TooManyRequests + + +@pytest.fixture +def mock_fastembed_embedding(): + """Mock FastEmbed TextEmbedding""" + mock = Mock() + # FastEmbed returns numpy arrays that need to be converted to lists + mock.embed.return_value = [ + np.array([0.1, 0.2, -0.3, 0.4, -0.5, 0.6, 0.7, -0.8, 0.9, -1.0]) + ] + return mock + + +@pytest.fixture +def processor_params(): + """Default parameters for FastEmbed processor""" + return { + "model": "test-embed-model", + "id": "test-fastembed", + "concurrency": 1 + } + + +@pytest.fixture +def mock_message(): + """Mock Pulsar message for FastEmbed processor""" + message = Mock() + message.properties.return_value = {"id": "test-msg-456"} + message.value.return_value = EmbeddingsRequest(text="FastEmbed test text") + return message + + +@pytest.fixture +def mock_flow(): + """Mock flow for EmbeddingsService testing""" + flow = Mock() + response_flow = Mock() + response_flow.send = AsyncMock() + flow.return_value = response_flow + flow.producer = {"response": Mock()} + flow.producer["response"].send = AsyncMock() + return flow + + +class TestFastEmbedProcessor: + """Test cases for FastEmbed embeddings processor""" + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') + def test_processor_initialization_default_params(self, mock_super_init, mock_text_embedding): + """Test processor initialization with default parameters""" + # Arrange + mock_embedding_instance = Mock() + mock_text_embedding.return_value = mock_embedding_instance + mock_super_init.return_value = None + + # Act + processor = Processor() + + # Assert + mock_text_embedding.assert_called_once_with( + model_name="sentence-transformers/all-MiniLM-L6-v2" + ) + assert processor.embeddings == mock_embedding_instance + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') + def test_processor_initialization_custom_model(self, mock_super_init, mock_text_embedding, processor_params): + """Test processor initialization with custom model""" + # Arrange + mock_embedding_instance = Mock() + mock_text_embedding.return_value = mock_embedding_instance + mock_super_init.return_value = None + + # Act + processor = Processor(**processor_params) + + # Assert + mock_text_embedding.assert_called_once_with( + model_name="test-embed-model" + ) + assert processor.embeddings == mock_embedding_instance + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_embeddings_successful(self, mock_text_embedding): + """Test successful embedding generation through on_embeddings method""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.return_value = [ + np.array([0.1, 0.2, -0.3, 0.4, -0.5]) + ] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Act + result = await processor.on_embeddings("Test text for embedding") + + # Assert + mock_embedding_instance.embed.assert_called_once_with(["Test text for embedding"]) + assert result == [[0.1, 0.2, -0.3, 0.4, -0.5]] + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_embeddings_multiple_vectors(self, mock_text_embedding): + """Test embedding generation returning multiple vectors""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.return_value = [ + np.array([0.1, 0.2, -0.3]), + np.array([0.4, -0.5, 0.6]) + ] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Act + result = await processor.on_embeddings("Multi-sentence text input") + + # Assert + assert result == [[0.1, 0.2, -0.3], [0.4, -0.5, 0.6]] + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_embeddings_empty_text(self, mock_text_embedding): + """Test embedding generation with empty text""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.return_value = [np.array([])] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Act + result = await processor.on_embeddings("") + + # Assert + mock_embedding_instance.embed.assert_called_once_with([""]) + assert result == [[]] + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_embeddings_large_dimensions(self, mock_text_embedding): + """Test embedding generation with large dimension vectors""" + # Arrange + mock_embedding_instance = Mock() + large_vector = np.random.rand(1536) # Common large embedding size + mock_embedding_instance.embed.return_value = [large_vector] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Act + result = await processor.on_embeddings("Text for large embedding") + + # Assert + assert len(result) == 1 + assert len(result[0]) == 1536 + assert result[0] == large_vector.tolist() + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_embeddings_unicode_text(self, mock_text_embedding): + """Test embedding generation with Unicode text""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.return_value = [ + np.array([0.3, -0.2, 0.1, 0.8, -0.6]) + ] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + unicode_text = "Hello 世界! 🚀 Café naïve résumé" + + # Act + result = await processor.on_embeddings(unicode_text) + + # Assert + mock_embedding_instance.embed.assert_called_once_with([unicode_text]) + assert result == [[0.3, -0.2, 0.1, 0.8, -0.6]] + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_embeddings_model_error(self, mock_text_embedding): + """Test handling of FastEmbed model errors""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.side_effect = Exception("Model loading failed") + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Act & Assert + with pytest.raises(Exception, match="Model loading failed"): + await processor.on_embeddings("Test text") + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_request_successful_flow(self, mock_text_embedding, mock_message, mock_flow): + """Test successful request handling through EmbeddingsService flow""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.return_value = [ + np.array([0.2, -0.1, 0.4, -0.3, 0.5]) + ] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + consumer = Mock() + + # Act + await processor.on_request(mock_message, consumer, mock_flow) + + # Assert + mock_embedding_instance.embed.assert_called_once_with(["FastEmbed test text"]) + + # Verify response sent through flow + mock_flow.assert_called_once_with("response") + mock_flow.return_value.send.assert_called_once() + + call_args = mock_flow.return_value.send.call_args + response = call_args[0][0] + properties = call_args[1]["properties"] + + assert isinstance(response, EmbeddingsResponse) + assert response.error is None + assert response.vectors == [[0.2, -0.1, 0.4, -0.3, 0.5]] + assert properties["id"] == "test-msg-456" + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_request_embedding_error(self, mock_text_embedding, mock_message, mock_flow): + """Test error handling in request processing""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.side_effect = Exception("FastEmbed error") + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + consumer = Mock() + + # Act + await processor.on_request(mock_message, consumer, mock_flow) + + # Assert + # Should send error response + mock_flow.producer["response"].send.assert_called_once() + + call_args = mock_flow.producer["response"].send.call_args + response = call_args[0][0] + properties = call_args[1]["properties"] + + assert isinstance(response, EmbeddingsResponse) + assert response.error is not None + assert response.error.type == "embeddings-error" + assert "FastEmbed error" in response.error.message + assert response.vectors is None + assert properties["id"] == "test-msg-456" + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_on_request_rate_limit_exception(self, mock_text_embedding, mock_message, mock_flow): + """Test handling of rate limit exceptions""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.side_effect = TooManyRequests("Rate limit exceeded") + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + consumer = Mock() + + # Act & Assert + with pytest.raises(TooManyRequests, match="Rate limit exceeded"): + await processor.on_request(mock_message, consumer, mock_flow) + + def test_add_args_method(self): + """Test that add_args method adds correct arguments""" + # Arrange + mock_parser = Mock() + + # Act + Processor.add_args(mock_parser) + + # Assert + # Verify model argument added + model_calls = [call for call in mock_parser.add_argument.call_args_list + if call[0][0] in ['-m', '--model']] + assert len(model_calls) == 1 + model_call = model_calls[0] + assert model_call[1]['default'] == "sentence-transformers/all-MiniLM-L6-v2" + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + def test_numpy_array_conversion(self, mock_text_embedding): + """Test that numpy arrays are properly converted to lists""" + # Arrange + mock_embedding_instance = Mock() + # Simulate FastEmbed returning numpy arrays with different dtypes + test_arrays = [ + np.array([0.1, 0.2, 0.3], dtype=np.float32), + np.array([-0.4, 0.5, -0.6], dtype=np.float64) + ] + mock_embedding_instance.embed.return_value = test_arrays + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Test conversion logic directly + result = [v.tolist() for v in test_arrays] + + # Assert + assert result == [[0.1, 0.2, 0.3], [-0.4, 0.5, -0.6]] + assert all(isinstance(vec, list) for vec in result) + assert all(isinstance(val, (int, float)) for vec in result for val in vec) + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + def test_processor_specifications_setup(self, mock_text_embedding): + """Test that processor sets up correct specifications""" + # Arrange + mock_text_embedding.return_value = Mock() + + # Act + processor = Processor() + + # Assert + # Verify processor inherits from EmbeddingsService correctly + assert hasattr(processor, 'specifications') + assert callable(processor.on_request) + assert hasattr(processor, 'on_embeddings') + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_batch_text_processing(self, mock_text_embedding): + """Test processing multiple texts in batch""" + # Arrange + mock_embedding_instance = Mock() + mock_embedding_instance.embed.return_value = [ + np.array([0.1, 0.2, 0.3]), + np.array([0.4, 0.5, 0.6]), + np.array([0.7, 0.8, 0.9]) + ] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Test multiple sequential calls (simulating batch processing) + texts = ["Text 1", "Text 2", "Text 3"] + results = [] + + # Act + for text in texts: + result = await processor.on_embeddings(text) + results.append(result) + + # Assert + assert len(results) == 3 + # Each call should process one text at a time + for i, call in enumerate(mock_embedding_instance.embed.call_args_list): + assert call[0][0] == [texts[i]] + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + def test_model_parameter_validation(self, mock_text_embedding): + """Test that model parameter is properly passed to FastEmbed""" + # Arrange + mock_text_embedding.return_value = Mock() + custom_model = "sentence-transformers/paraphrase-MiniLM-L6-v2" + + # Act + processor = Processor(model=custom_model) + + # Assert + mock_text_embedding.assert_called_once_with(model_name=custom_model) + + @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') + async def test_dimension_consistency(self, mock_text_embedding): + """Test that embedding dimensions are consistent across calls""" + # Arrange + mock_embedding_instance = Mock() + # Always return same dimension vectors + consistent_dimension = 384 + mock_embedding_instance.embed.side_effect = [ + [np.random.rand(consistent_dimension)], + [np.random.rand(consistent_dimension)], + [np.random.rand(consistent_dimension)] + ] + mock_text_embedding.return_value = mock_embedding_instance + + processor = Processor() + + # Act + results = [] + for i in range(3): + result = await processor.on_embeddings(f"Test text {i}") + results.append(result) + + # Assert + for result in results: + assert len(result) == 1 # One embedding per call + assert len(result[0]) == consistent_dimension # Consistent dimensions \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_ollama_core.py b/tests/unit/test_embeddings/test_ollama_core.py new file mode 100644 index 00000000..5cab8b6a --- /dev/null +++ b/tests/unit/test_embeddings/test_ollama_core.py @@ -0,0 +1,222 @@ +""" +Unit tests for Ollama core embedding functionality + +Tests the core business logic without full processor initialization, +focusing on the embedding generation methods and client interaction. +""" + +import pytest +from unittest.mock import patch, Mock +import os + + +class TestOllamaCore: + """Test core Ollama embedding functionality""" + + @patch('trustgraph.embeddings.ollama.processor.Client') + def test_ollama_client_initialization(self, mock_client_class): + """Test Ollama client initialization""" + # Arrange + mock_client = Mock() + mock_client_class.return_value = mock_client + + # Act + from ollama import Client + client = Client(host="http://localhost:11434") + + # Assert + mock_client_class.assert_called_once_with(host="http://localhost:11434") + + @patch('trustgraph.embeddings.ollama.processor.Client') + def test_ollama_embed_method(self, mock_client_class): + """Test Ollama client embed method functionality""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock( + embeddings=[0.1, 0.2, -0.3, 0.4, -0.5] + ) + mock_client_class.return_value = mock_client + + # Act + from ollama import Client + client = Client(host="http://localhost:11434") + result = client.embed(model="test-model", input="Test text") + + # Assert + mock_client.embed.assert_called_once_with(model="test-model", input="Test text") + assert result.embeddings == [0.1, 0.2, -0.3, 0.4, -0.5] + + def test_environment_variable_handling(self): + """Test OLLAMA_HOST environment variable handling""" + # Test default value + from trustgraph.embeddings.ollama.processor import default_ollama + + # Should use localhost by default + assert default_ollama == os.getenv("OLLAMA_HOST", 'http://localhost:11434') + + # Test with custom environment variable + with patch.dict(os.environ, {'OLLAMA_HOST': 'http://custom:8080'}): + import importlib + # Reload module to pick up new env var + custom_host = os.getenv("OLLAMA_HOST", 'http://localhost:11434') + assert custom_host == 'http://custom:8080' + + def test_default_model_constant(self): + """Test that default model constant is properly defined""" + # Act + from trustgraph.embeddings.ollama.processor import default_model + + # Assert + assert default_model == "mxbai-embed-large" + assert isinstance(default_model, str) + assert len(default_model) > 0 + + def test_embedding_response_structure(self): + """Test Ollama embedding response structure""" + # Arrange + mock_response = Mock() + mock_response.embeddings = [0.1, 0.2, 0.3, 0.4, 0.5] + + # Act + embeddings = mock_response.embeddings + + # Assert + assert isinstance(embeddings, list) + assert len(embeddings) == 5 + assert all(isinstance(val, (int, float)) for val in embeddings) + + def test_embedding_dimension_consistency(self): + """Test that embeddings maintain consistent dimensions""" + # Arrange + test_responses = [ + Mock(embeddings=[0.1, 0.2, 0.3, 0.4, 0.5]), + Mock(embeddings=[0.6, 0.7, 0.8, 0.9, 1.0]), + Mock(embeddings=[-0.1, -0.2, -0.3, -0.4, -0.5]) + ] + + # Act + embeddings = [resp.embeddings for resp in test_responses] + + # Assert + dimensions = [len(emb) for emb in embeddings] + assert all(dim == dimensions[0] for dim in dimensions), "All embeddings should have same dimension" + assert dimensions[0] == 5 + + def test_text_input_handling(self): + """Test various text input scenarios""" + # Arrange + test_cases = [ + "", # Empty string + "Simple text", # Normal text + "Text with 123 numbers", # Alphanumeric + "Hello 世界! 🌍", # Unicode + "A" * 10000, # Very long text + "Multi\nline\ntext", # Multiline + ] + + # Act & Assert + for text in test_cases: + # Should all be valid string inputs + assert isinstance(text, str) + # Text length should be accessible + assert len(text) >= 0 + + def test_error_handling_simulation(self): + """Test error handling in Ollama client simulation""" + # Arrange + mock_client = Mock() + mock_client.embed.side_effect = Exception("Connection failed") + + # Act & Assert + with pytest.raises(Exception, match="Connection failed"): + mock_client.embed(model="test-model", input="Test text") + + def test_model_not_found_simulation(self): + """Test handling when model is not found simulation""" + # Arrange + mock_client = Mock() + mock_client.embed.side_effect = Exception("Model 'unknown-model' not found") + + # Act & Assert + with pytest.raises(Exception, match="Model 'unknown-model' not found"): + mock_client.embed(model="unknown-model", input="Test text") + + def test_concurrent_requests_simulation(self): + """Test concurrent request handling simulation""" + # Arrange + import asyncio + + async def mock_embed_call(text_id): + # Simulate async embedding call + await asyncio.sleep(0.001) # Minimal delay + return [float(text_id)] * 5 + + # Act + async def run_concurrent_test(): + tasks = [mock_embed_call(i) for i in range(5)] + results = await asyncio.gather(*tasks) + return results + + # Run the test + import asyncio + results = asyncio.run(run_concurrent_test()) + + # Assert + assert len(results) == 5 + for i, result in enumerate(results): + assert result == [float(i)] * 5 + + def test_processor_add_args_functionality(self): + """Test add_args static method functionality""" + # Arrange + from trustgraph.embeddings.ollama.processor import Processor + mock_parser = Mock() + + # Act + Processor.add_args(mock_parser) + + # Assert + # Verify model argument was added + model_calls = [call for call in mock_parser.add_argument.call_args_list + if call[0][0] in ['-m', '--model']] + assert len(model_calls) == 1 + model_call = model_calls[0] + assert model_call[1]['default'] == "mxbai-embed-large" + + # Verify ollama argument was added + ollama_calls = [call for call in mock_parser.add_argument.call_args_list + if call[0][0] in ['-r', '--ollama']] + assert len(ollama_calls) == 1 + ollama_call = ollama_calls[0] + assert ollama_call[1]['default'] == 'http://localhost:11434' + + def test_host_url_validation(self): + """Test Ollama host URL validation""" + # Arrange + valid_hosts = [ + "http://localhost:11434", + "http://127.0.0.1:11434", + "http://ollama-server:11434", + "https://secure-ollama.com:443" + ] + + # Act & Assert + for host in valid_hosts: + assert isinstance(host, str) + assert host.startswith(("http://", "https://")) + assert ":" in host # Should have port specification + + def test_embedding_vector_properties(self): + """Test properties of embedding vectors""" + # Arrange + test_embedding = [0.1, -0.2, 0.3, -0.4, 0.5, 0.0, -0.1, 0.8] + + # Act + positive_count = sum(1 for x in test_embedding if x > 0) + negative_count = sum(1 for x in test_embedding if x < 0) + zero_count = sum(1 for x in test_embedding if x == 0) + + # Assert + assert positive_count + negative_count + zero_count == len(test_embedding) + assert all(isinstance(x, (int, float)) for x in test_embedding) + assert all(-1.0 <= x <= 1.0 for x in test_embedding) # Typical embedding range \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_ollama_processor.py b/tests/unit/test_embeddings/test_ollama_processor.py new file mode 100644 index 00000000..3cf03c91 --- /dev/null +++ b/tests/unit/test_embeddings/test_ollama_processor.py @@ -0,0 +1,334 @@ +""" +Unit tests for Ollama embeddings processor + +Tests the core business logic of Ollama embedding generation without +relying on external Ollama service infrastructure. +""" + +import pytest +from unittest.mock import patch, Mock, AsyncMock, MagicMock +import os + +from trustgraph.embeddings.ollama.processor import Processor +from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error +from trustgraph.exceptions import TooManyRequests + + +@pytest.fixture +def mock_ollama_client(): + """Mock Ollama client for testing""" + client = Mock() + client.embed.return_value = Mock( + embeddings=[0.1, 0.2, -0.3, 0.4, -0.5, 0.6, 0.7, -0.8, 0.9, -1.0] + ) + return client + + +@pytest.fixture +def processor_params(): + """Default parameters for Ollama processor""" + return { + "model": "test-embed-model", + "ollama": "http://localhost:11434", + "input_queue": "test-input", + "output_queue": "test-output", + "subscriber": "test-subscriber" + } + + +@pytest.fixture +def mock_message(): + """Mock Pulsar message""" + message = Mock() + message.properties.return_value = {"id": "test-msg-123"} + message.value.return_value = EmbeddingsRequest(text="Test embedding text") + return message + + +class TestOllamaEmbeddingsProcessor: + """Test cases for Ollama embeddings processor""" + + @patch('trustgraph.embeddings.ollama.processor.Client') + def test_processor_initialization_default_params(self, mock_client_class): + """Test processor initialization with default parameters""" + # Arrange + mock_client = Mock() + mock_client_class.return_value = mock_client + + # Act + processor = Processor() + + # Assert + assert processor.model == "mxbai-embed-large" # default model + mock_client_class.assert_called_once_with(host='http://localhost:11434') + assert processor.client == mock_client + + @patch('trustgraph.embeddings.ollama.processor.Client') + def test_processor_initialization_custom_params(self, mock_client_class, processor_params): + """Test processor initialization with custom parameters""" + # Arrange + mock_client = Mock() + mock_client_class.return_value = mock_client + + # Act + processor = Processor(**processor_params) + + # Assert + assert processor.model == "test-embed-model" + mock_client_class.assert_called_once_with(host="http://localhost:11434") + assert processor.client == mock_client + + @patch('trustgraph.embeddings.ollama.processor.Client') + def test_processor_initialization_env_variable(self, mock_client_class): + """Test processor uses OLLAMA_HOST environment variable""" + # Arrange + mock_client = Mock() + mock_client_class.return_value = mock_client + + with patch.dict(os.environ, {'OLLAMA_HOST': 'http://custom-host:8080'}): + # Act + processor = Processor() + + # Assert + mock_client_class.assert_called_once_with(host='http://custom-host:8080') + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_handle_successful_embedding(self, mock_client_class, mock_message): + """Test successful embedding generation""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock( + embeddings=[0.1, 0.2, -0.3, 0.4, -0.5] + ) + mock_client_class.return_value = mock_client + + processor = Processor(model="test-model") + processor.send = AsyncMock() + + # Act + await processor.handle(mock_message) + + # Assert + mock_client.embed.assert_called_once_with( + model="test-model", + input="Test embedding text" + ) + + # Verify response sent + processor.send.assert_called_once() + call_args = processor.send.call_args + response = call_args[0][0] + properties = call_args[1]["properties"] + + assert isinstance(response, EmbeddingsResponse) + assert response.error is None + assert response.vectors == [0.1, 0.2, -0.3, 0.4, -0.5] + assert properties["id"] == "test-msg-123" + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_handle_client_error(self, mock_client_class, mock_message): + """Test handling of Ollama client errors""" + # Arrange + mock_client = Mock() + mock_client.embed.side_effect = Exception("Ollama connection failed") + mock_client_class.return_value = mock_client + + processor = Processor() + processor.send = AsyncMock() + + # Act & Assert + with pytest.raises(Exception, match="Ollama connection failed"): + await processor.handle(mock_message) + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_handle_model_not_found(self, mock_client_class, mock_message): + """Test handling when Ollama model is not found""" + # Arrange + mock_client = Mock() + mock_client.embed.side_effect = Exception("Model 'unknown-model' not found") + mock_client_class.return_value = mock_client + + processor = Processor(model="unknown-model") + processor.send = AsyncMock() + + # Act & Assert + with pytest.raises(Exception, match="Model 'unknown-model' not found"): + await processor.handle(mock_message) + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_handle_empty_text(self, mock_client_class): + """Test handling of empty text input""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock(embeddings=[]) + mock_client_class.return_value = mock_client + + processor = Processor() + processor.send = AsyncMock() + + empty_message = Mock() + empty_message.properties.return_value = {"id": "empty-msg"} + empty_message.value.return_value = EmbeddingsRequest(text="") + + # Act + await processor.handle(empty_message) + + # Assert + mock_client.embed.assert_called_once_with( + model="mxbai-embed-large", + input="" + ) + + processor.send.assert_called_once() + response = processor.send.call_args[0][0] + assert response.vectors == [] + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_handle_large_text_input(self, mock_client_class): + """Test handling of large text input""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock( + embeddings=[0.1] * 1536 # Large embedding vector + ) + mock_client_class.return_value = mock_client + + processor = Processor() + processor.send = AsyncMock() + + large_text = "A" * 10000 # Large text input + large_message = Mock() + large_message.properties.return_value = {"id": "large-msg"} + large_message.value.return_value = EmbeddingsRequest(text=large_text) + + # Act + await processor.handle(large_message) + + # Assert + mock_client.embed.assert_called_once_with( + model="mxbai-embed-large", + input=large_text + ) + + processor.send.assert_called_once() + response = processor.send.call_args[0][0] + assert len(response.vectors) == 1536 + + @patch('trustgraph.embeddings.ollama.processor.Client') + def test_embedding_vector_consistency(self, mock_client_class): + """Test that embedding vectors have consistent dimensions""" + # Arrange + mock_client = Mock() + mock_client_class.return_value = mock_client + + processor = Processor(model="test-model") + + # Assert + assert processor.model == "test-model" + assert processor.client == mock_client + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_handle_unicode_text(self, mock_client_class): + """Test handling of Unicode text input""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock( + embeddings=[0.2, -0.1, 0.3, -0.4, 0.5] + ) + mock_client_class.return_value = mock_client + + processor = Processor() + processor.send = AsyncMock() + + unicode_text = "Hello 世界! 🌍 Émoji test" + unicode_message = Mock() + unicode_message.properties.return_value = {"id": "unicode-msg"} + unicode_message.value.return_value = EmbeddingsRequest(text=unicode_text) + + # Act + await processor.handle(unicode_message) + + # Assert + mock_client.embed.assert_called_once_with( + model="mxbai-embed-large", + input=unicode_text + ) + + processor.send.assert_called_once() + response = processor.send.call_args[0][0] + assert response.vectors == [0.2, -0.1, 0.3, -0.4, 0.5] + + def test_add_args_method(self): + """Test that add_args method adds correct arguments""" + # Arrange + mock_parser = Mock() + + # Act + Processor.add_args(mock_parser) + + # Assert + # Verify model argument added + model_calls = [call for call in mock_parser.add_argument.call_args_list + if call[0][0] in ['-m', '--model']] + assert len(model_calls) == 1 + model_call = model_calls[0] + assert model_call[1]['default'] == "mxbai-embed-large" + + # Verify ollama argument added + ollama_calls = [call for call in mock_parser.add_argument.call_args_list + if call[0][0] in ['-r', '--ollama']] + assert len(ollama_calls) == 1 + ollama_call = ollama_calls[0] + assert ollama_call[1]['default'] == 'http://localhost:11434' + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_message_properties_handling(self, mock_client_class, mock_message): + """Test proper handling of message properties""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock(embeddings=[0.1, 0.2, 0.3]) + mock_client_class.return_value = mock_client + + processor = Processor() + processor.send = AsyncMock() + + # Act + await processor.handle(mock_message) + + # Assert + mock_message.properties.assert_called_once() + + # Verify properties passed to send + call_properties = processor.send.call_args[1]["properties"] + assert call_properties["id"] == "test-msg-123" + + @patch('trustgraph.embeddings.ollama.processor.Client') + async def test_concurrent_requests_handling(self, mock_client_class): + """Test that processor can handle multiple concurrent requests""" + # Arrange + mock_client = Mock() + mock_client.embed.return_value = Mock(embeddings=[0.1, 0.2, 0.3]) + mock_client_class.return_value = mock_client + + processor = Processor() + processor.send = AsyncMock() + + # Create multiple messages + messages = [] + for i in range(3): + msg = Mock() + msg.properties.return_value = {"id": f"msg-{i}"} + msg.value.return_value = EmbeddingsRequest(text=f"Text {i}") + messages.append(msg) + + # Act - Handle messages concurrently + import asyncio + await asyncio.gather(*[processor.handle(msg) for msg in messages]) + + # Assert + assert mock_client.embed.call_count == 3 + assert processor.send.call_count == 3 + + # Verify each message was handled correctly + for i, call in enumerate(mock_client.embed.call_args_list): + assert call[1]["input"] == f"Text {i}" \ No newline at end of file