Spaces:
Runtime error
Runtime error
File size: 2,664 Bytes
9fb584d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
"""
Tests for the pharmaceutical data management agent workflow.
"""
import unittest
import os
import json
from unittest.mock import MagicMock, patch
from anthropic import Anthropic
from langgraph.graph import START
from agents.state import AgentState
from graph.workflow import create_agent_graph
class TestAgentWorkflow(unittest.TestCase):
"""Test case for the agent workflow."""
def setUp(self):
"""Set up test fixtures."""
# Mock the Anthropic client
self.anthropic_client = MagicMock()
self.anthropic_client.messages.create.return_value = MagicMock(
content=[MagicMock(text="INTENT_COMPLETE\nThis is a test intent.")]
)
# Mock the database
self.db = MagicMock()
self.db.get_tables.return_value = {
"raw_tables": ["RAW_SALES_TRANSACTIONS", "RAW_HCP_DATA", "RAW_PRODUCT_DATA"],
"staging_tables": ["STG_SALES", "STG_HCP", "STG_PRODUCT"],
"ard_tables": ["ARD_SALES_PERFORMANCE", "ARD_HCP_ENGAGEMENT", "ARD_MARKET_ANALYSIS"],
"data_products": ["DP_SALES_DASHBOARD", "DP_HCP_TARGETING"]
}
# Create a test agent graph
self.agent_graph, self.update_state_dict = create_agent_graph(
self.anthropic_client, self.db
)
def test_workflow_has_entry_point(self):
"""Test that the workflow has an entry point."""
# Get the graph builder
graph_builder = self.agent_graph.builder
# Check if there's an edge from START
start_edges = [edge for edge in graph_builder.edges if edge[0] == START]
# Assert that there's at least one edge from START
self.assertTrue(len(start_edges) > 0, "Workflow must have at least one edge from START")
def test_workflow_execution(self):
"""Test the workflow execution."""
# Prepare the input state
input_state = {
"messages": [{"role": "user", "content": "Create a dashboard showing sales by region"}],
"user_intent": {},
"data_context": {},
"pipeline_plan": {},
"sql_queries": [],
"execution_results": {},
"confidence_scores": {},
"status": "planning",
"current_agent": "understanding_agent"
}
# Invoke the workflow
result = self.agent_graph.invoke(input_state)
# Check that the workflow updated the state
self.assertIn("user_intent", result)
self.assertTrue(result["user_intent"].get("understood", False))
if __name__ == "__main__":
unittest.main() |