74 lines
3.3 KiB
Python
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())
|