gateway/tests/unit/datamodels/test_workflow_models.py

#!/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,
    TaskResult
)
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"])