diff --git a/ocr/tests/test_gemini_engine.py b/ocr/tests/test_gemini_engine.py index bf709e4..3674b4a 100644 --- a/ocr/tests/test_gemini_engine.py +++ b/ocr/tests/test_gemini_engine.py @@ -2,11 +2,11 @@ Covers: GeminiEngine initialization, PDF size validation, successful extraction, empty results, and error handling. -All Vertex AI SDK calls are mocked. +All google-genai SDK calls are mocked. """ import json -from unittest.mock import MagicMock, patch, PropertyMock +from unittest.mock import MagicMock, patch import pytest @@ -156,22 +156,16 @@ class TestExtractMaintenance: }, ] - mock_model = MagicMock() - mock_model.generate_content.return_value = _make_gemini_response(schedule) + mock_client = MagicMock() + mock_client.models.generate_content.return_value = _make_gemini_response(schedule) - with ( - patch( - "app.engines.gemini_engine.importlib_vertex_ai" - ) if False else patch.dict("sys.modules", { - "google.cloud": MagicMock(), - "google.cloud.aiplatform": MagicMock(), - "vertexai": MagicMock(), - "vertexai.generative_models": MagicMock(), - }), - ): + with patch.dict("sys.modules", { + "google.genai": MagicMock(), + "google.genai.types": MagicMock(), + }): engine = GeminiEngine() - engine._model = mock_model - engine._generation_config = MagicMock() + engine._client = mock_client + engine._model_name = "gemini-2.5-flash" result = engine.extract_maintenance(_make_pdf_bytes()) @@ -200,12 +194,12 @@ class TestExtractMaintenance: mock_settings.vertex_ai_location = "us-central1" mock_settings.gemini_model = "gemini-2.5-flash" - mock_model = MagicMock() - mock_model.generate_content.return_value = _make_gemini_response([]) + mock_client = MagicMock() + mock_client.models.generate_content.return_value = _make_gemini_response([]) engine = GeminiEngine() - engine._model = mock_model - engine._generation_config = MagicMock() + engine._client = mock_client + engine._model_name = "gemini-2.5-flash" result = engine.extract_maintenance(_make_pdf_bytes()) @@ -223,12 +217,12 @@ class TestExtractMaintenance: schedule = [{"serviceName": "Brake Fluid Replacement"}] - mock_model = MagicMock() - mock_model.generate_content.return_value = _make_gemini_response(schedule) + mock_client = MagicMock() + mock_client.models.generate_content.return_value = _make_gemini_response(schedule) engine = GeminiEngine() - engine._model = mock_model - engine._generation_config = MagicMock() + engine._client = mock_client + engine._model_name = "gemini-2.5-flash" result = engine.extract_maintenance(_make_pdf_bytes()) @@ -264,7 +258,8 @@ class TestErrorHandling: with ( patch("app.engines.gemini_engine.settings") as mock_settings, patch.dict("sys.modules", { - "google.cloud.aiplatform": None, + "google": None, + "google.genai": None, }), ): mock_settings.google_vision_key_path = "/fake/creds.json" @@ -283,12 +278,12 @@ class TestErrorHandling: mock_settings.vertex_ai_location = "us-central1" mock_settings.gemini_model = "gemini-2.5-flash" - mock_model = MagicMock() - mock_model.generate_content.side_effect = RuntimeError("API quota exceeded") + mock_client = MagicMock() + mock_client.models.generate_content.side_effect = RuntimeError("API quota exceeded") engine = GeminiEngine() - engine._model = mock_model - engine._generation_config = MagicMock() + engine._client = mock_client + engine._model_name = "gemini-2.5-flash" with pytest.raises(GeminiProcessingError, match="maintenance extraction failed"): engine.extract_maintenance(_make_pdf_bytes()) @@ -307,12 +302,12 @@ class TestErrorHandling: mock_response = MagicMock() mock_response.text = "not valid json {{" - mock_model = MagicMock() - mock_model.generate_content.return_value = mock_response + mock_client = MagicMock() + mock_client.models.generate_content.return_value = mock_response engine = GeminiEngine() - engine._model = mock_model - engine._generation_config = MagicMock() + engine._client = mock_client + engine._model_name = "gemini-2.5-flash" with pytest.raises(GeminiProcessingError, match="invalid JSON"): engine.extract_maintenance(_make_pdf_bytes()) @@ -322,32 +317,32 @@ class TestErrorHandling: class TestLazyInitialization: - """Verify the model is not created until first use.""" + """Verify the client is not created until first use.""" - def test_model_is_none_after_construction(self): - """GeminiEngine should not initialize the model in __init__.""" + def test_client_is_none_after_construction(self): + """GeminiEngine should not initialize the client in __init__.""" engine = GeminiEngine() - assert engine._model is None + assert engine._client is None @patch("app.engines.gemini_engine.settings") @patch("app.engines.gemini_engine.os.path.isfile", return_value=True) - def test_model_reused_on_second_call(self, mock_isfile, mock_settings): - """Once initialized, the same model instance is reused.""" + def test_client_reused_on_second_call(self, mock_isfile, mock_settings): + """Once initialized, the same client instance is reused.""" mock_settings.google_vision_key_path = "/fake/creds.json" mock_settings.vertex_ai_project = "test-project" mock_settings.vertex_ai_location = "us-central1" mock_settings.gemini_model = "gemini-2.5-flash" schedule = [{"serviceName": "Oil Change", "intervalMiles": 5000}] - mock_model = MagicMock() - mock_model.generate_content.return_value = _make_gemini_response(schedule) + mock_client = MagicMock() + mock_client.models.generate_content.return_value = _make_gemini_response(schedule) engine = GeminiEngine() - engine._model = mock_model - engine._generation_config = MagicMock() + engine._client = mock_client + engine._model_name = "gemini-2.5-flash" engine.extract_maintenance(_make_pdf_bytes()) engine.extract_maintenance(_make_pdf_bytes()) - # Model's generate_content should have been called twice - assert mock_model.generate_content.call_count == 2 + # Client's generate_content should have been called twice + assert mock_client.models.generate_content.call_count == 2