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