| """ |
| Unit tests for LangGraph multi-agent orchestrator. |
| """ |
| from __future__ import annotations |
|
|
| import pytest |
| from unittest.mock import patch, AsyncMock, MagicMock |
| from httpx import AsyncClient |
|
|
| from agents.planner import planner_node |
| from agents.graph import run_copilot_pipeline |
| from app.services.llm_gateway import LLMResponse |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_planner_node_classification(): |
| mock_llm_res = LLMResponse( |
| content='{"routing_decision": "retrieval", "explanation": "Asks for policy", "doc_type_filter": "hr"}', |
| provider="groq", |
| model="llama-3.1-8b-instant", |
| tokens_used=20, |
| latency_ms=100.0, |
| ) |
|
|
| with patch("app.services.llm_gateway.llm_gateway.generate", new_callable=AsyncMock) as mock_gen: |
| mock_gen.return_value = mock_llm_res |
|
|
| initial_state = { |
| "query": "What is our leave policy?", |
| "user_role": "employee", |
| } |
| res_state = await planner_node(initial_state) |
|
|
| assert res_state["routing_decision"] == "retrieval" |
| assert res_state["doc_type_filter"] == "hr" |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_copilot_pipeline_execution(): |
| mock_llm_res = LLMResponse( |
| content="According to the policy, employees receive 20 days of paid leave per year.", |
| provider="groq", |
| model="llama-3.1-8b-instant", |
| tokens_used=45, |
| latency_ms=150.0, |
| ) |
|
|
| with patch("app.services.llm_gateway.llm_gateway.generate", new_callable=AsyncMock) as mock_gen: |
| mock_gen.return_value = mock_llm_res |
|
|
| result = await run_copilot_pipeline( |
| query="Hello copilot", |
| tenant_id="tenant-123", |
| user_id="user-456", |
| ) |
|
|
| assert "response" in result |
| assert result["latency_ms"] >= 0 |
| assert "routing_decision" in result |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_chat_api_endpoint(client: AsyncClient): |
| |
| reg = await client.post("/api/v1/auth/register", json={ |
| "email": "agent_user@company.com", |
| "password": "Password123!", |
| "full_name": "Agent User", |
| "tenant_name": "Agent Org" |
| }) |
| token = reg.json()["tokens"]["access_token"] |
| headers = {"Authorization": f"Bearer {token}"} |
|
|
| |
| mock_llm_res = LLMResponse( |
| content="Hello! How can I assist you with your enterprise automation today?", |
| provider="groq", |
| model="llama-3.1-8b-instant", |
| tokens_used=12, |
| latency_ms=50.0, |
| ) |
|
|
| with patch("app.services.llm_gateway.llm_gateway.generate", new_callable=AsyncMock) as mock_gen: |
| mock_gen.return_value = mock_llm_res |
|
|
| chat_req = {"message": "Hi there!"} |
| response = await client.post("/api/v1/chat/", headers=headers, json=chat_req) |
|
|
| assert response.status_code == 200, response.text |
| data = response.json() |
| assert "response" in data |
| assert "session_id" in data |
| assert "routing_decision" in data |
|
|
