146 lines
5.4 KiB
Python
146 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Unit tests for AI service (mainServiceAi.py)
|
|
Tests callAiContent, callAiPlanning, and related functionality.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, AsyncMock, patch
|
|
|
|
from modules.datamodels.datamodelAi import AiCallOptions, OperationTypeEnum, PriorityEnum, ProcessingModeEnum
|
|
from modules.datamodels.datamodelExtraction import ContentPart
|
|
from modules.datamodels.datamodelWorkflow import AiResponse
|
|
|
|
|
|
class TestAiServiceCallAiContent:
|
|
"""Test callAiContent method (mocked)"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callAiContent_requires_operationType(self):
|
|
"""Test that callAiContent requires operationType to be set"""
|
|
from modules.services.serviceAi.mainServiceAi import AiService
|
|
|
|
# Create mock services
|
|
mockServices = Mock()
|
|
mockServices.workflow = None
|
|
mockServices.chat = Mock()
|
|
mockServices.chat.progressLogStart = Mock()
|
|
mockServices.chat.progressLogUpdate = Mock()
|
|
mockServices.chat.progressLogFinish = Mock()
|
|
mockServices.chat.storeWorkflowStat = Mock()
|
|
|
|
aiService = AiService(mockServices)
|
|
|
|
# Mock aiObjects initialization
|
|
aiService.aiObjects = Mock()
|
|
aiService._ensureAiObjectsInitialized = AsyncMock()
|
|
|
|
# Test with missing operationType - should analyze prompt
|
|
options = AiCallOptions() # operationType not set
|
|
options.operationType = None
|
|
|
|
# Mock _analyzePromptAndCreateOptions
|
|
analyzedOptions = AiCallOptions()
|
|
analyzedOptions.operationType = OperationTypeEnum.DATA_ANALYSE
|
|
aiService._analyzePromptAndCreateOptions = AsyncMock(return_value=analyzedOptions)
|
|
|
|
# Mock _callAiWithLooping
|
|
aiService._callAiWithLooping = AsyncMock(return_value="Test response")
|
|
|
|
# Mock aiObjects.call
|
|
mockResponse = Mock()
|
|
mockResponse.content = "Test response"
|
|
aiService.aiObjects.call = AsyncMock(return_value=mockResponse)
|
|
|
|
# Call should work (will analyze prompt if operationType not set)
|
|
result = await aiService.callAiContent(
|
|
prompt="Test prompt",
|
|
options=options
|
|
)
|
|
|
|
# Should have analyzed prompt and set operationType
|
|
assert result is not None
|
|
assert isinstance(result, AiResponse)
|
|
|
|
|
|
class TestAiServiceCallAiPlanning:
|
|
"""Test callAiPlanning method (mocked)"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callAiPlanning_basic(self):
|
|
"""Test basic callAiPlanning call"""
|
|
from modules.services.serviceAi.mainServiceAi import AiService
|
|
|
|
# Create mock services
|
|
mockServices = Mock()
|
|
mockServices.workflow = None
|
|
mockServices.utils = Mock()
|
|
mockServices.utils.writeDebugFile = Mock()
|
|
|
|
aiService = AiService(mockServices)
|
|
|
|
# Mock aiObjects
|
|
aiService.aiObjects = Mock()
|
|
mockResponse = Mock()
|
|
mockResponse.content = '{"result": "plan"}'
|
|
aiService.aiObjects.call = AsyncMock(return_value=mockResponse)
|
|
aiService._ensureAiObjectsInitialized = AsyncMock()
|
|
|
|
# Call planning
|
|
result = await aiService.callAiPlanning(
|
|
prompt="Test planning prompt"
|
|
)
|
|
|
|
assert result == '{"result": "plan"}'
|
|
|
|
|
|
class TestAiServiceOperationTypeHandling:
|
|
"""Test operationType handling in callAiContent"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callAiContent_with_outputFormat_sets_documentGenerate(self):
|
|
"""Test that outputFormat sets operationType to DOCUMENT_GENERATE"""
|
|
from modules.services.serviceAi.mainServiceAi import AiService
|
|
|
|
mockServices = Mock()
|
|
mockServices.workflow = None
|
|
mockServices.chat = Mock()
|
|
mockServices.chat.progressLogStart = Mock()
|
|
mockServices.chat.progressLogUpdate = Mock()
|
|
mockServices.chat.progressLogFinish = Mock()
|
|
mockServices.utils = Mock()
|
|
mockServices.utils.jsonExtractString = Mock(return_value='{"documents": []}')
|
|
|
|
aiService = AiService(mockServices)
|
|
aiService.aiObjects = Mock()
|
|
aiService._ensureAiObjectsInitialized = AsyncMock()
|
|
|
|
# Mock _callAiWithLooping
|
|
aiService._callAiWithLooping = AsyncMock(return_value='{"documents": []}')
|
|
|
|
# Mock generation service
|
|
with patch('modules.services.serviceGeneration.mainServiceGeneration.GenerationService') as mockGenService:
|
|
mockGenInstance = Mock()
|
|
mockGenInstance.renderReport = AsyncMock(return_value=(b"content", "application/pdf"))
|
|
mockGenService.return_value = mockGenInstance
|
|
|
|
options = AiCallOptions() # operationType not set
|
|
options.operationType = None
|
|
|
|
# Should set operationType to DOCUMENT_GENERATE when outputFormat is provided
|
|
try:
|
|
result = await aiService.callAiContent(
|
|
prompt="Generate document",
|
|
options=options,
|
|
outputFormat="pdf"
|
|
)
|
|
# If it gets here, operationType was set correctly
|
|
assert options.operationType == OperationTypeEnum.DOCUMENT_GENERATE
|
|
except Exception:
|
|
# If it fails, that's okay for unit test - we're testing the logic
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|