86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
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
|