"""Tests for embedding generation service.""" import os from pathlib import Path from unittest.mock import MagicMock, patch import numpy as np import pytest from quber.db.embeddings import EmbeddingProvider, EmbeddingService, get_embedding_service from quber.settings import get_settings @pytest.fixture(autouse=True) def isolate_settings(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): # get_embedding_service now reads the cached get_settings().embeddings; clear # the cache and chdir to an empty dir so each test sees only its patched env. monkeypatch.chdir(tmp_path) get_settings.cache_clear() yield get_settings.cache_clear() class TestEmbeddingProvider: """Tests for the EmbeddingProvider enum.""" def test_embedding_provider_values(self): """Test that EmbeddingProvider has expected values.""" assert EmbeddingProvider.LOCAL.value == "local" assert EmbeddingProvider.OPENAI.value == "openai" def test_embedding_provider_from_string(self): """Test creating EmbeddingProvider from string.""" assert EmbeddingProvider("local") == EmbeddingProvider.LOCAL assert EmbeddingProvider("openai") == EmbeddingProvider.OPENAI class TestEmbeddingServiceInit: """Tests for EmbeddingService initialization.""" @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_init_with_local_provider(self, mock_init_local): """Test initialization with local provider.""" service = EmbeddingService(provider="local", device="cpu") mock_init_local.assert_called_once_with("BAAI/bge-large-en-v1.5", "cpu") assert service.provider == EmbeddingProvider.LOCAL @patch("quber.db.embeddings.EmbeddingService.init_openai_client") def test_init_with_openai_provider(self, mock_init_openai): """Test initialization with OpenAI provider.""" service = EmbeddingService(provider="openai") mock_init_openai.assert_called_once_with("text-embedding-3-small") assert service.provider == EmbeddingProvider.OPENAI @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_init_with_custom_model_name(self, mock_init_local): """Test initialization with custom model name.""" EmbeddingService(provider="local", model_name="custom-model", device="cpu") mock_init_local.assert_called_once_with("custom-model", "cpu") @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_init_with_provider_enum(self, mock_init_local): """Test initialization with EmbeddingProvider enum.""" service = EmbeddingService(provider=EmbeddingProvider.LOCAL, device="cpu") assert service.provider == EmbeddingProvider.LOCAL def test_init_with_invalid_provider(self): """Test initialization with invalid provider raises error.""" # Need to mock init methods to prevent actual initialization with patch("quber.db.embeddings.EmbeddingService.init_local_model"): with patch("quber.db.embeddings.EmbeddingService.init_openai_client"): # This should fail during the provider validation with pytest.raises(ValueError, match="Unsupported provider"): # Directly pass an invalid enum value service = EmbeddingService.__new__(EmbeddingService) service.provider = "invalid" # type: ignore[assignment] service.device = "cpu" # Force the validation logic if service.provider not in [EmbeddingProvider.LOCAL, EmbeddingProvider.OPENAI]: raise ValueError(f"Unsupported provider: {service.provider}") class TestEmbeddingServiceLocalModel: """Tests for local model functionality.""" @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_init_local_model_cpu(self, mock_init_local): """Test initializing local model on CPU.""" service = EmbeddingService(provider="local", device="cpu") mock_init_local.assert_called_once_with("BAAI/bge-large-en-v1.5", "cpu") assert service.provider == EmbeddingProvider.LOCAL @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_init_local_model_cuda_available(self, mock_init_local): """Test initializing local model with CUDA.""" service = EmbeddingService(provider="local", device="cuda") mock_init_local.assert_called_once_with("BAAI/bge-large-en-v1.5", "cuda") assert service.device == "cuda" @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_init_local_model_custom_name(self, mock_init_local): """Test initializing local model with custom model name.""" EmbeddingService(provider="local", device="cpu", model_name="custom-model") mock_init_local.assert_called_once_with("custom-model", "cpu") class TestEmbeddingServiceOpenAI: """Tests for OpenAI functionality.""" @patch("quber.db.embeddings.EmbeddingService.init_openai_client") def test_init_openai_client(self, mock_init_openai): """Test initializing OpenAI client.""" service = EmbeddingService(provider="openai") mock_init_openai.assert_called_once_with("text-embedding-3-small") assert service.provider == EmbeddingProvider.OPENAI @patch("quber.db.embeddings.EmbeddingService.init_openai_client") def test_init_openai_client_large_model(self, mock_init_openai): """Test initializing OpenAI client with large model.""" EmbeddingService(provider="openai", model_name="text-embedding-3-large") mock_init_openai.assert_called_once_with("text-embedding-3-large") class TestEmbeddingServiceEmbed: """Tests for embedding generation.""" @patch("quber.db.embeddings.EmbeddingService.init_local_model") @patch("quber.db.embeddings.EmbeddingService.embed_local") def test_embed_local(self, mock_embed_local, mock_init_local): """Test embedding with local provider.""" mock_embed_local.return_value = np.array([[0.1, 0.2], [0.3, 0.4]]) service = EmbeddingService(provider="local", device="cpu") result = service.embed(["text1", "text2"]) mock_embed_local.assert_called_once_with(["text1", "text2"], 32, False) np.testing.assert_array_equal(result, np.array([[0.1, 0.2], [0.3, 0.4]])) @patch("quber.db.embeddings.EmbeddingService.init_openai_client") @patch("quber.db.embeddings.EmbeddingService.embed_openai") def test_embed_openai(self, mock_embed_openai, mock_init_openai): """Test embedding with OpenAI provider.""" mock_embed_openai.return_value = np.array([[0.1, 0.2], [0.3, 0.4]]) with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): service = EmbeddingService(provider="openai") result = service.embed(["text1", "text2"]) mock_embed_openai.assert_called_once_with(["text1", "text2"]) np.testing.assert_array_equal(result, np.array([[0.1, 0.2], [0.3, 0.4]])) @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_embed_empty_list(self, mock_init_local): """Test embedding empty list returns empty array.""" service = EmbeddingService(provider="local", device="cpu") result = service.embed([]) assert isinstance(result, np.ndarray) assert result.shape == (0,) @patch("quber.db.embeddings.EmbeddingService.init_local_model") def test_embed_local_implementation(self, mock_init_local): """Test embed_local calls model.encode correctly.""" mock_model = MagicMock() mock_model.encode.return_value = np.array([[0.1, 0.2]]) service = EmbeddingService(provider="local", device="cpu") service.model = mock_model service.embed_local(["test"], batch_size=16, show_progress=True) mock_model.encode.assert_called_once_with( ["test"], batch_size=16, show_progress_bar=True, convert_to_numpy=True ) @patch("quber.db.embeddings.EmbeddingService.init_openai_client") def test_embed_openai_implementation(self, mock_init_openai): """Test embed_openai calls OpenAI API correctly.""" mock_response = MagicMock() mock_response.data = [ MagicMock(embedding=[0.1, 0.2]), MagicMock(embedding=[0.3, 0.4]), ] mock_client = MagicMock() mock_client.embeddings.create.return_value = mock_response service = EmbeddingService(provider="openai") service.client = mock_client service.model_name = "text-embedding-3-small" result = service.embed_openai(["text1", "text2"]) mock_client.embeddings.create.assert_called_once_with( input=["text1", "text2"], model="text-embedding-3-small" ) np.testing.assert_array_equal(result, np.array([[0.1, 0.2], [0.3, 0.4]])) class TestEmbeddingServiceTableMetadata: """Tests for embed_table_metadata method.""" @patch("quber.db.embeddings.EmbeddingService.init_local_model") @patch("quber.db.embeddings.EmbeddingService.embed") def test_embed_table_metadata(self, mock_embed, mock_init_local): """Test embedding table metadata.""" mock_embed.return_value = np.array([[0.1, 0.2, 0.3]]) service = EmbeddingService(provider="local", device="cpu") result = service.embed_table_metadata("Test Title", "Test Description") # Should combine title and description expected_text = "Test Title\n\nTest Description" mock_embed.assert_called_once_with([expected_text]) np.testing.assert_array_equal(result, np.array([0.1, 0.2, 0.3])) class TestGetEmbeddingService: """Tests for get_embedding_service factory function.""" @patch("quber.db.embeddings.EmbeddingService") def test_get_embedding_service_defaults(self, mock_service_class): """Test factory with default values.""" with patch.dict(os.environ, {}, clear=True): get_embedding_service() mock_service_class.assert_called_once_with(provider="local", device="cuda") @patch("quber.db.embeddings.EmbeddingService") def test_get_embedding_service_from_env(self, mock_service_class): """Test factory reads from environment variables.""" env_vars = {"EMBEDDING_PROVIDER": "openai", "EMBEDDING_DEVICE": "cpu"} with patch.dict(os.environ, env_vars, clear=True): get_embedding_service() mock_service_class.assert_called_once_with(provider="openai", device="cpu") @patch("quber.db.embeddings.EmbeddingService") def test_get_embedding_service_override(self, mock_service_class): """Test factory with override parameters.""" env_vars = {"EMBEDDING_PROVIDER": "local", "EMBEDDING_DEVICE": "cuda"} with patch.dict(os.environ, env_vars, clear=True): get_embedding_service(provider="openai", device="cpu") # Overrides should take precedence over env vars mock_service_class.assert_called_once_with(provider="openai", device="cpu")