test: add AI extraction pipeline tests (mocked)
This commit is contained in:
85
backend/tests/test_ai_extraction.py
Normal file
85
backend/tests/test_ai_extraction.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user