Files
ea-chatbot-lg/backend/tests/test_query_analyzer_refinement.py
Yunxiao Xu 68c0985482 feat(auth): Complete OIDC security refactor and modernize test suite
- 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.
2026-02-15 02:50:26 -08:00

103 lines
3.9 KiB
Python

import pytest
from unittest.mock import MagicMock, patch
from langchain_core.messages import HumanMessage, AIMessage
from ea_chatbot.graph.nodes.query_analyzer import query_analyzer_node, QueryAnalysis
@pytest.fixture
def base_state():
return {
"messages": [],
"question": "",
"analysis": None,
"next_action": "",
"summary": ""
}
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
def test_refinement_coreference_from_history(mock_get_llm, base_state):
"""
Test that the analyzer can resolve Year/State from history.
User asks "What about in NJ?" after a Florida 2024 query.
Expected: next_action = 'plan', NOT 'clarify' due to missing year.
"""
state = base_state.copy()
state["messages"] = [
HumanMessage(content="Show me 2024 results for Florida"),
AIMessage(content="Here are the 2024 results for Florida...")
]
state["question"] = "What about in New Jersey?"
state["summary"] = "The user is looking for 2024 election results."
mock_llm = MagicMock()
mock_get_llm.return_value = mock_llm
mock_structured = MagicMock()
mock_llm.with_structured_output.return_value = mock_structured
# We expect the LLM to eventually return 'plan' because it sees the context.
# For now, if it returns 'clarify', this test should fail once we update the prompt to BE less strict.
mock_structured.invoke.return_value = QueryAnalysis(
data_required=["2024 results", "New Jersey"],
unknowns=[],
ambiguities=[],
conditions=["state=NJ", "year=2024"],
next_action="plan"
)
result = query_analyzer_node(state)
assert result["next_action"] == "plan"
assert "NJ" in str(result["analysis"]["conditions"])
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
def test_refinement_tolerance_for_missing_format(mock_get_llm, base_state):
"""
Test that the analyzer doesn't flag missing output format or database name.
User asks "Give me a graph of turnout".
Expected: next_action = 'plan', even if 'format' or 'db' is not in query.
"""
state = base_state.copy()
state["question"] = "Give me a graph of voter turnout in 2024 for Florida"
mock_llm = MagicMock()
mock_get_llm.return_value = mock_llm
mock_structured = MagicMock()
mock_llm.with_structured_output.return_value = mock_structured
mock_structured.invoke.return_value = QueryAnalysis(
data_required=["voter turnout", "Florida"],
unknowns=[],
ambiguities=[],
conditions=["year=2024"],
next_action="plan"
)
result = query_analyzer_node(state)
assert result["next_action"] == "plan"
# Ensure no ambiguities were added by the analyzer itself (hallucinated requirement)
assert len(result["analysis"]["ambiguities"]) == 0
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
def test_refinement_enforces_voter_identity_clarification(mock_get_llm, base_state):
"""
Test that 'track the same voter' still triggers clarification.
"""
state = base_state.copy()
state["question"] = "Track the same voter participation in 2020 and 2024."
mock_llm = MagicMock()
mock_get_llm.return_value = mock_llm
mock_structured = MagicMock()
mock_llm.with_structured_output.return_value = mock_structured
# We WANT it to clarify here because voter identity is not defined.
mock_structured.invoke.return_value = QueryAnalysis(
data_required=["voter participation"],
unknowns=[],
ambiguities=["Please define what fields constitute 'the same voter' (e.g. ID, or Name and DOB)."],
conditions=[],
next_action="clarify"
)
result = query_analyzer_node(state)
assert result["next_action"] == "clarify"
assert "identity" in str(result["analysis"]["ambiguities"]).lower() or "same voter" in str(result["analysis"]["ambiguities"]).lower()