agents-runtime/test_orchestrator.py

74 lines
3.3 KiB
Python

import asyncio
from unittest.mock import patch, MagicMock
from temporalio.client import Client
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from src.orchestrator import CentralOrchestratorWorkflow
from src.agent import run_agent_activity
# Mock requests.post to avoid real network requests during testing
def mock_post(url, *args, **kwargs):
mock_resp = MagicMock()
if "utility-agent" in url:
mock_resp.status_code = 200
mock_resp.text = "Mocked utility scraping complete: 15 bills processed."
return mock_resp
elif "whatsapp-gateway" in url:
mock_resp.status_code = 200
mock_resp.json.return_value = {"status": "sent"}
return mock_resp
mock_resp.status_code = 404
return mock_resp
async def main():
with patch("src.agent.requests.post", side_effect=mock_post) as mock_requests_post:
async with await WorkflowEnvironment.start_time_skipping() as env:
async with Worker(
env.client,
task_queue="agent-task-queue",
workflows=[CentralOrchestratorWorkflow],
activities=[run_agent_activity],
):
# Test Case 1: Reminder Workflow (Original test modified to fix payload mismatch)
print("Submitting Workflow 1: Reminder & WhatsApp request...")
payload_reminder = {
"text": "Remind me to call John at 4pm and send me a WhatsApp about it",
"targetJid": "whatsapp-group-123",
"sender": "Alice"
}
result_reminder = await env.client.execute_workflow(
CentralOrchestratorWorkflow.run,
payload_reminder,
id="test-workflow-reminder",
task_queue="agent-task-queue",
)
print(f"Workflow 1 (Reminder) Result: {result_reminder}")
# Test Case 2: Scrape Utility Data Workflow (New test case)
print("\nSubmitting Workflow 2: Scrape utility data request...")
payload_scraper = {
"text": "Please scrape the utility data from the portal",
"targetJid": "whatsapp-group-123",
"sender": "Alice"
}
result_scraper = await env.client.execute_workflow(
CentralOrchestratorWorkflow.run,
payload_scraper,
id="test-workflow-scraper",
task_queue="agent-task-queue",
)
print(f"Workflow 2 (Scraper) Result: {result_scraper}")
# Verify that the trigger_utility_scraper tool was called
called_urls = [call.args[0] for call in mock_requests_post.call_args_list]
print(f"\nCaptured HTTP POST URLs: {called_urls}")
utility_url = "http://utility-agent.ai-agents.svc.cluster.local:5005/api/run"
if utility_url in called_urls:
print("SUCCESS: trigger_utility_scraper tool was successfully called by the LLM!")
else:
print("FAILURE: trigger_utility_scraper tool was NOT called by the LLM.")
if __name__ == "__main__":
asyncio.run(main())