#!/usr/bin/env python3 # Copyright (c) 2025 Patrick Motsch # All rights reserved. """ Unit tests for workflow models in datamodelWorkflow.py Tests ActionDefinition, AiResponse, ExtractContentParameters, and workflow-level models. """ import pytest import json from typing import Dict, Any from modules.datamodels.datamodelWorkflow import ( ActionDefinition, AiResponse, AiResponseMetadata, DocumentData, ExtractContentParameters, RequestContext, UnderstandingResult, TaskDefinition, WorkflowTaskResult ) from modules.datamodels.datamodelDocref import DocumentReferenceList, DocumentListReference from modules.datamodels.datamodelAi import OperationTypeEnum from modules.datamodels.datamodelExtraction import ExtractionOptions, MergeStrategy class TestActionDefinition: """Test ActionDefinition model""" def test_actionDefinition_creation(self): """Test creating ActionDefinition with required fields""" actionDef = ActionDefinition( action="ai.process", actionObjective="Process documents with AI" ) assert actionDef.action == "ai.process" assert actionDef.actionObjective == "Process documents with AI" assert actionDef.parameters is None assert actionDef.documentList is None assert actionDef.connectionReference is None def test_actionDefinition_needsStage2_without_parameters(self): """Test needsStage2() returns True when parameters are None""" actionDef = ActionDefinition( action="ai.process", actionObjective="Process documents" ) assert actionDef.needsStage2() is True def test_actionDefinition_needsStage2_with_parameters(self): """Test needsStage2() returns False when parameters are set""" actionDef = ActionDefinition( action="ai.process", actionObjective="Process documents", parameters={"resultType": "pdf"} ) assert actionDef.needsStage2() is False def test_actionDefinition_hasParameters(self): """Test hasParameters() method""" actionDef = ActionDefinition( action="ai.process", actionObjective="Process documents" ) assert actionDef.hasParameters() is False actionDef.parameters = {"resultType": "pdf"} assert actionDef.hasParameters() is True def test_actionDefinition_with_documentList(self): """Test ActionDefinition with documentList""" docList = DocumentReferenceList(references=[ DocumentListReference(label="task1_results") ]) actionDef = ActionDefinition( action="ai.process", actionObjective="Process documents", documentList=docList ) assert actionDef.documentList is not None assert len(actionDef.documentList.references) == 1 class TestAiResponse: """Test AiResponse model""" def test_aiResponse_creation(self): """Test creating AiResponse with content""" response = AiResponse(content='{"result": "success"}') assert response.content == '{"result": "success"}' assert response.metadata is None assert response.documents is None def test_aiResponse_with_metadata(self): """Test AiResponse with metadata""" metadata = AiResponseMetadata( title="Test Document", operationType="dataGenerate" ) response = AiResponse( content='{"result": "success"}', metadata=metadata ) assert response.metadata.title == "Test Document" assert response.metadata.operationType == "dataGenerate" def test_aiResponse_with_documents(self): """Test AiResponse with documents""" doc = DocumentData( documentName="test.pdf", documentData=b"PDF content", mimeType="application/pdf" ) response = AiResponse( content='{"result": "success"}', documents=[doc] ) assert len(response.documents) == 1 assert response.documents[0].documentName == "test.pdf" def test_aiResponse_toJson_valid_json(self): """Test toJson() with valid JSON content""" response = AiResponse(content='{"result": "success", "data": [1, 2, 3]}') result = response.toJson() assert isinstance(result, dict) assert result["result"] == "success" assert result["data"] == [1, 2, 3] def test_aiResponse_toJson_list_wrapped(self): """Test toJson() wraps list in dict""" response = AiResponse(content='[1, 2, 3]') result = response.toJson() assert isinstance(result, dict) assert "data" in result assert result["data"] == [1, 2, 3] class TestExtractContentParameters: """Test ExtractContentParameters model""" def test_extractContentParameters_creation(self): """Test creating ExtractContentParameters""" docList = DocumentReferenceList(references=[ DocumentListReference(label="test_docs") ]) params = ExtractContentParameters(documentList=docList) assert params.documentList is not None assert params.extractionOptions is None def test_extractContentParameters_with_options(self): """Test ExtractContentParameters with extractionOptions""" docList = DocumentReferenceList(references=[ DocumentListReference(label="test_docs") ]) mergeStrategy = MergeStrategy( mergeType="concatenate", groupBy="typeGroup" ) options = ExtractionOptions( prompt="Extract all content", mergeStrategy=mergeStrategy ) params = ExtractContentParameters( documentList=docList, extractionOptions=options ) assert params.extractionOptions is not None assert params.extractionOptions.prompt == "Extract all content" class TestDocumentData: """Test DocumentData model""" def test_documentData_creation(self): """Test creating DocumentData""" doc = DocumentData( documentName="test.txt", documentData="Test content", mimeType="text/plain" ) assert doc.documentName == "test.txt" assert doc.documentData == "Test content" assert doc.mimeType == "text/plain" def test_documentData_with_bytes(self): """Test DocumentData with bytes data""" doc = DocumentData( documentName="test.pdf", documentData=b"PDF bytes", mimeType="application/pdf" ) assert isinstance(doc.documentData, bytes) class TestRequestContext: """Test RequestContext model""" def test_requestContext_creation(self): """Test creating RequestContext""" context = RequestContext( originalPrompt="Test prompt", userLanguage="en", detectedComplexity="simple" ) assert context.originalPrompt == "Test prompt" assert context.userLanguage == "en" assert context.detectedComplexity == "simple" assert context.requiresDocuments is False assert context.requiresWebResearch is False class TestTaskDefinition: """Test TaskDefinition model""" def test_taskDefinition_creation(self): """Test creating TaskDefinition""" task = TaskDefinition( id="task1", objective="Complete task", deliverable={"type": "document", "format": "pdf"} ) assert task.id == "task1" assert task.objective == "Complete task" assert task.requiresContentGeneration is True assert task.requiresWebResearch is False if __name__ == "__main__": pytest.main([__file__, "-v"])