Testing embeedings

This commit is contained in:
Cyber MacGeddon 2025-07-14 15:55:12 +01:00
parent 44e14dff8b
commit f7437cf72d
8 changed files with 1923 additions and 0 deletions

View file

@ -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
"""

View file

@ -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"
}

View file

@ -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

View file

@ -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]

View file

@ -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"

View file

@ -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

View file

@ -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

View file

@ -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}"