157 lines
5.8 KiB
Python
157 lines
5.8 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
"""
|
|
Integration tests for workflow execution
|
|
Tests full workflow execution with state management, Stage 1/2, document extraction flow.
|
|
"""
|
|
|
|
import pytest
|
|
import uuid
|
|
from unittest.mock import Mock, AsyncMock, patch
|
|
|
|
from modules.datamodels.datamodelChat import ChatWorkflow, TaskContext, TaskStep
|
|
from modules.datamodels.datamodelWorkflow import ActionDefinition
|
|
from modules.datamodels.datamodelDocref import DocumentReferenceList, DocumentListReference, DocumentItemReference
|
|
|
|
|
|
class TestWorkflowStateManagement:
|
|
"""Test workflow state management during execution"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_state_increments(self):
|
|
"""Test that workflow state increments correctly during execution"""
|
|
workflow = ChatWorkflow(
|
|
id=str(uuid.uuid4()),
|
|
name="Test Workflow",
|
|
mandateId="test_mandate"
|
|
)
|
|
|
|
# Initial state
|
|
assert workflow.currentRound == 0
|
|
assert workflow.currentTask == 0
|
|
assert workflow.currentAction == 0
|
|
|
|
# Simulate workflow progression
|
|
workflow.incrementAction()
|
|
assert workflow.currentAction == 1
|
|
|
|
workflow.incrementTask()
|
|
assert workflow.currentTask == 1
|
|
assert workflow.currentAction == 0 # Reset when task increments
|
|
|
|
workflow.incrementRound()
|
|
assert workflow.currentRound == 1
|
|
assert workflow.currentTask == 0 # Reset when round increments
|
|
assert workflow.currentAction == 0
|
|
|
|
|
|
class TestStage1ToStage2Flow:
|
|
"""Test Stage 1 → Stage 2 parameter generation flow"""
|
|
|
|
def test_actionDefinition_needsStage2_logic(self):
|
|
"""Test needsStage2() deterministic logic"""
|
|
# Stage 1: No parameters
|
|
actionDef = ActionDefinition(
|
|
action="ai.process",
|
|
actionObjective="Process documents"
|
|
)
|
|
assert actionDef.needsStage2() is True
|
|
|
|
# Stage 2: Parameters added
|
|
actionDef.parameters = {"resultType": "pdf"}
|
|
assert actionDef.needsStage2() is False
|
|
|
|
def test_actionDefinition_stage1_resources(self):
|
|
"""Test that Stage 1 always defines documentList and connectionReference if needed"""
|
|
docList = DocumentReferenceList(references=[
|
|
DocumentListReference(label="task1_results")
|
|
])
|
|
actionDef = ActionDefinition(
|
|
action="ai.process",
|
|
actionObjective="Process documents",
|
|
documentList=docList,
|
|
connectionReference="conn123"
|
|
)
|
|
# Stage 1 resources are set, but parameters are not
|
|
assert actionDef.documentList is not None
|
|
assert actionDef.connectionReference == "conn123"
|
|
assert actionDef.needsStage2() is True # Still needs Stage 2 for parameters
|
|
|
|
|
|
class TestDocumentExtractionFlow:
|
|
"""Test document extraction → AI processing flow"""
|
|
|
|
def test_extractContentParameters_structure(self):
|
|
"""Test ExtractContentParameters structure"""
|
|
from modules.datamodels.datamodelWorkflow import ExtractContentParameters
|
|
|
|
docList = DocumentReferenceList(references=[
|
|
DocumentListReference(label="input_docs")
|
|
])
|
|
params = ExtractContentParameters(documentList=docList)
|
|
|
|
assert params.documentList is not None
|
|
assert len(params.documentList.references) == 1
|
|
assert params.extractionOptions is None # Optional
|
|
|
|
def test_documentReferenceList_parsing(self):
|
|
"""Test DocumentReferenceList parsing from strings"""
|
|
stringList = [
|
|
"docList:task1_results",
|
|
"docItem:doc123:test.pdf"
|
|
]
|
|
refList = DocumentReferenceList.from_string_list(stringList)
|
|
|
|
assert len(refList.references) == 2
|
|
assert isinstance(refList.references[0], DocumentListReference)
|
|
assert isinstance(refList.references[1], DocumentItemReference)
|
|
|
|
|
|
class TestDocumentReferenceLookup:
|
|
"""Test document reference lookup across tasks/rounds"""
|
|
|
|
def test_documentListReference_with_messageId(self):
|
|
"""Test DocumentListReference with messageId for cross-round references"""
|
|
ref = DocumentListReference(
|
|
messageId="msg123",
|
|
label="task1_results"
|
|
)
|
|
assert ref.messageId == "msg123"
|
|
assert ref.label == "task1_results"
|
|
assert ref.to_string() == "docList:msg123:task1_results"
|
|
|
|
def test_documentListReference_without_messageId(self):
|
|
"""Test DocumentListReference without messageId (current message)"""
|
|
ref = DocumentListReference(label="task1_results")
|
|
assert ref.messageId is None
|
|
assert ref.to_string() == "docList:task1_results"
|
|
|
|
|
|
class TestJsonParsing:
|
|
"""Test JSON parsing with broken/incomplete JSON"""
|
|
|
|
def test_parseJsonWithModel_with_code_fences(self):
|
|
"""Test parseJsonWithModel handles code fences"""
|
|
from modules.shared.jsonUtils import parseJsonWithModel
|
|
|
|
jsonStr = '```json\n{"action": "ai.process", "actionObjective": "Process"}\n```'
|
|
result = parseJsonWithModel(jsonStr, ActionDefinition)
|
|
|
|
assert isinstance(result, ActionDefinition)
|
|
assert result.action == "ai.process"
|
|
|
|
def test_parseJsonWithModel_with_extra_text(self):
|
|
"""Test parseJsonWithModel extracts JSON from text with extra content"""
|
|
from modules.shared.jsonUtils import parseJsonWithModel
|
|
|
|
jsonStr = 'Some text before {"action": "ai.process", "actionObjective": "Process"} some text after'
|
|
result = parseJsonWithModel(jsonStr, ActionDefinition)
|
|
|
|
assert isinstance(result, ActionDefinition)
|
|
assert result.action == "ai.process"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|