Spaces:
Sleeping
Sleeping
File size: 4,933 Bytes
db4ba8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
Tests for AI node error handling
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.ai.nodes.extract import llm_extraction_node
from src.ai.state import ExtractionGraphState
@pytest.mark.asyncio
async def test_extraction_node_handles_missing_doc_fields():
"""Test that extraction node validates required document fields."""
state = ExtractionGraphState(
batch_id="test-batch",
documents=[
{
"doc_id": "doc-1",
# Missing 'pages' field
"storage_path": "s3://bucket/doc.pdf"
}
]
)
result = await llm_extraction_node(state)
# Should mark document as error instead of crashing
assert len(result["documents"]) == 1
assert "error" in result["documents"][0]
assert result["documents"][0].get("fallback_required") is True
@pytest.mark.asyncio
async def test_extraction_node_handles_empty_documents():
"""Test that extraction node handles empty document list."""
state = ExtractionGraphState(
batch_id="test-batch",
documents=[]
)
result = await llm_extraction_node(state)
# Should return gracefully with no documents
assert result["documents"] == []
assert result["combined_data"] == {}
@pytest.mark.asyncio
async def test_extraction_node_specific_exception_handling():
"""Test that extraction node catches only specific exceptions."""
state = ExtractionGraphState(
batch_id="test-batch",
documents=[
{
"doc_id": "doc-1",
"pages": ["base64-encoded-image"],
"storage_path": "s3://bucket/doc.pdf"
}
]
)
# Mock LLM to raise ValueError (expected)
with patch("src.ai.nodes.extract.ChatGoogleGenerativeAI") as mock_llm:
mock_instance = MagicMock()
structured_llm = AsyncMock()
structured_llm.ainvoke = AsyncMock(side_effect=ValueError("Malformed input"))
mock_instance.with_structured_output.return_value = structured_llm
mock_llm.return_value = mock_instance
result = await llm_extraction_node(state)
# Should handle ValueError gracefully
assert len(result["documents"]) == 1
assert "error" in result["documents"][0]
@pytest.mark.asyncio
async def test_extraction_node_reraises_unknown_exceptions():
"""Test that extraction node re-raises unexpected exceptions."""
state = ExtractionGraphState(
batch_id="test-batch",
documents=[
{
"doc_id": "doc-1",
"pages": ["base64-encoded-image"],
"storage_path": "s3://bucket/doc.pdf"
}
]
)
# Mock LLM to raise unexpected exception
with patch("src.ai.nodes.extract.ChatGoogleGenerativeAI") as mock_llm:
mock_instance = MagicMock()
structured_llm = AsyncMock()
structured_llm.ainvoke = AsyncMock(side_effect=RuntimeError("Unexpected API error"))
mock_instance.with_structured_output.return_value = structured_llm
mock_llm.return_value = mock_instance
# Should re-raise the unexpected exception
with pytest.raises(RuntimeError, match="Unexpected API error"):
await llm_extraction_node(state)
@pytest.mark.asyncio
async def test_extraction_node_combines_data_correctly():
"""Test that extraction node correctly combines data from multiple docs."""
state = ExtractionGraphState(
batch_id="test-batch",
documents=[
{
"doc_id": "doc-1",
"pages": ["page1"],
"storage_path": "s3://bucket/doc1.pdf"
},
{
"doc_id": "doc-2",
"pages": ["page2"],
"storage_path": "s3://bucket/doc2.pdf"
}
]
)
# Mock LLM responses
with patch("src.ai.nodes.extract.ChatGoogleGenerativeAI") as mock_llm:
mock_instance = MagicMock()
structured_llm = AsyncMock()
# Return different data for each document
responses = [
MagicMock(model_dump=MagicMock(return_value={"importer_name": "Company A", "cif_value": 1000})),
MagicMock(model_dump=MagicMock(return_value={"importer_name": "Company B", "cif_value": 2000}))
]
structured_llm.ainvoke = AsyncMock(side_effect=responses)
mock_instance.with_structured_output.return_value = structured_llm
mock_llm.return_value = mock_instance
result = await llm_extraction_node(state)
# Combined data should have the last writer's value
assert result["combined_data"]["importer_name"] == "Company B"
assert result["combined_data"]["cif_value"] == 2000
|