From 436a3cdd97e7b13121794332945d1b9f613acb90 Mon Sep 17 00:00:00 2001
From: Daniel Bedeleanu <bede@tfm.ro>
Date: Sat, 18 Apr 2026 17:02:41 +0000
Subject: [PATCH] test: add AI extraction pipeline tests (mocked)

---
 backend/tests/test_ai_extraction.py | 85 +++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)
 create mode 100644 backend/tests/test_ai_extraction.py

diff --git a/backend/tests/test_ai_extraction.py b/backend/tests/test_ai_extraction.py
new file mode 100644
index 00000000..59844478
--- /dev/null
+++ b/backend/tests/test_ai_extraction.py
@@ -0,0 +1,85 @@
+import pytest
+from fastapi import status
+from unittest.mock import patch
+
+
+class TestAIExtraction:
+    """Test AI label extraction pipeline (mocked)."""
+
+    def test_gemini_extraction(self, test_client, mock_gemini):
+        """Test AI extraction using Gemini."""
+        response = test_client.post(
+            "/api/ai/extract",
+            json={
+                "image_base64": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
+                "provider": "gemini"
+            }
+        )
+        assert response.status_code == status.HTTP_200_OK
+        data = response.json()
+        assert "name" in data
+        assert data["name"] == "Test Item"
+        assert data["part_number"] == "PN-12345"
+
+    def test_claude_extraction(self, test_client, mock_claude):
+        """Test AI extraction using Claude."""
+        response = test_client.post(
+            "/api/ai/extract",
+            json={
+                "image_base64": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
+                "provider": "claude"
+            }
+        )
+        assert response.status_code == status.HTTP_200_OK
+        data = response.json()
+        assert data["name"] == "Test Item"
+
+    def test_extraction_box_mode(self, test_client, mock_gemini):
+        """Test AI extraction in 'box' mode (focus on labels)."""
+        response = test_client.post(
+            "/api/ai/extract",
+            json={
+                "image_base64": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
+                "provider": "gemini",
+                "mode": "box"
+            }
+        )
+        assert response.status_code == status.HTTP_200_OK
+
+    def test_extraction_invalid_provider(self, test_client):
+        """Test that invalid provider fails."""
+        response = test_client.post(
+            "/api/ai/extract",
+            json={
+                "image_base64": "invalid",
+                "provider": "invalid_provider"
+            }
+        )
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
+
+    def test_extraction_missing_image(self, test_client):
+        """Test that missing image fails."""
+        response = test_client.post(
+            "/api/ai/extract",
+            json={"provider": "gemini"}
+        )
+        assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+class TestAIValidation:
+    """Test AI extraction validation (user confirmation before save)."""
+
+    def test_validate_extraction(self, test_client, test_db):
+        """Test that extracted data requires user validation."""
+        response = test_client.post(
+            "/api/ai/extract",
+            json={
+                "image_base64": "xyz",
+                "provider": "gemini",
+                "save": False
+            }
+        )
+        assert response.status_code == status.HTTP_200_OK
+        data = response.json()
+        assert "extracted_data" in data
+        assert "user_confirmation_required" in data or data.get("save") == False