68c0985482
- Refactored OIDC flow to implement PKCE, state/nonce validation, and BFF pattern. - Centralized configuration in Settings class (DEV_MODE, FRONTEND_URL, OIDC_REDIRECT_URI). - Updated auth routers to use conditional secure cookie flags based on DEV_MODE. - Modernized and cleaned up test suite by removing legacy Streamlit tests. - Fixed linting errors and unused imports across the backend.
92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
from unittest.mock import MagicMock, patch
|
|
from ea_chatbot.graph.workflow import app
|
|
from ea_chatbot.graph.nodes.query_analyzer import QueryAnalysis
|
|
from ea_chatbot.schemas import TaskPlanResponse, TaskPlanContext, CodeGenerationResponse
|
|
|
|
from langchain_core.messages import AIMessage
|
|
|
|
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
|
|
@patch("ea_chatbot.graph.nodes.planner.get_llm_model")
|
|
@patch("ea_chatbot.graph.nodes.coder.get_llm_model")
|
|
@patch("ea_chatbot.graph.nodes.summarizer.get_llm_model")
|
|
@patch("ea_chatbot.graph.nodes.researcher.get_llm_model")
|
|
@patch("ea_chatbot.utils.database_inspection.get_data_summary")
|
|
@patch("ea_chatbot.graph.nodes.executor.Settings")
|
|
@patch("ea_chatbot.graph.nodes.executor.DBClient")
|
|
def test_workflow_full_flow(mock_db_client, mock_settings, mock_get_summary, mock_researcher_llm, mock_summarizer_llm, mock_coder_llm, mock_planner_llm, mock_qa_llm):
|
|
"""Test the flow from query_analyzer through planner to coder."""
|
|
|
|
# Mock Settings for Executor
|
|
mock_settings_instance = MagicMock()
|
|
mock_settings_instance.db_host = "localhost"
|
|
mock_settings_instance.db_port = 5432
|
|
mock_settings_instance.db_user = "user"
|
|
mock_settings_instance.db_pswd = "pass"
|
|
mock_settings_instance.db_name = "test_db"
|
|
mock_settings_instance.db_table = "test_table"
|
|
mock_settings.return_value = mock_settings_instance
|
|
|
|
# Mock DBClient
|
|
mock_client_instance = MagicMock()
|
|
mock_db_client.return_value = mock_client_instance
|
|
|
|
# 1. Mock Query Analyzer
|
|
mock_qa_instance = MagicMock()
|
|
mock_qa_llm.return_value = mock_qa_instance
|
|
mock_qa_instance.with_structured_output.return_value.invoke.return_value = QueryAnalysis(
|
|
data_required=["2024 results"],
|
|
unknowns=[],
|
|
ambiguities=[],
|
|
conditions=[],
|
|
next_action="plan"
|
|
)
|
|
|
|
# 2. Mock Planner
|
|
mock_planner_instance = MagicMock()
|
|
mock_planner_llm.return_value = mock_planner_instance
|
|
mock_get_summary.return_value = "Data summary"
|
|
mock_planner_instance.with_structured_output.return_value.invoke.return_value = TaskPlanResponse(
|
|
goal="Task Goal",
|
|
reflection="Reflection",
|
|
context=TaskPlanContext(initial_context="Ctx", assumptions=[], constraints=[]),
|
|
steps=["Step 1"]
|
|
)
|
|
|
|
# 3. Mock Coder
|
|
mock_coder_instance = MagicMock()
|
|
mock_coder_llm.return_value = mock_coder_instance
|
|
mock_coder_instance.with_structured_output.return_value.invoke.return_value = CodeGenerationResponse(
|
|
code="print('Hello')",
|
|
explanation="Explanation"
|
|
)
|
|
|
|
# 4. Mock Summarizer
|
|
mock_summarizer_instance = MagicMock()
|
|
mock_summarizer_llm.return_value = mock_summarizer_instance
|
|
mock_summarizer_instance.invoke.return_value = AIMessage(content="Summary")
|
|
|
|
# 5. Mock Researcher (just in case)
|
|
mock_researcher_instance = MagicMock()
|
|
mock_researcher_llm.return_value = mock_researcher_instance
|
|
|
|
# Initial state
|
|
initial_state = {
|
|
"messages": [],
|
|
"question": "Show me the 2024 results",
|
|
"analysis": None,
|
|
"next_action": "",
|
|
"plan": None,
|
|
"code": None,
|
|
"error": None,
|
|
"plots": [],
|
|
"dfs": {}
|
|
}
|
|
|
|
# Run the graph
|
|
# We use recursion_limit to avoid infinite loops in placeholders if any
|
|
result = app.invoke(initial_state, config={"recursion_limit": 10})
|
|
|
|
assert result["next_action"] == "plan"
|
|
assert "plan" in result and result["plan"] is not None
|
|
assert "code" in result and "print('Hello')" in result["code"]
|
|
assert "analysis" in result |