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