""" tests/unit/test_agent_nodes.py Tests for individual LangGraph nodes with mocked LLM/DB calls. """ import json import pytest from unittest.mock import patch, MagicMock # ── Intent Router ────────────────────────────────────────────────────────────── @pytest.mark.unit class TestIntentRouter: def test_routes_sql_intent(self, sample_state, mocker): mocker.patch( "agent.nodes.intent_router.get_groq_client" ).return_value.complete_system.return_value = json.dumps( {"intent": "sql", "reasoning": "needs aggregation"} ) from agent.nodes.intent_router import intent_router result = intent_router(sample_state) assert result["intent"] == "sql" def test_routes_pandas_intent(self, sample_state, mocker): mocker.patch( "agent.nodes.intent_router.get_groq_client" ).return_value.complete_system.return_value = json.dumps( {"intent": "pandas", "reasoning": "csv manipulation"} ) from agent.nodes.intent_router import intent_router result = intent_router({**sample_state, "connector_id": "csv:http://fake"}) assert result["intent"] == "pandas" def test_routes_unsupported_intent(self, sample_state, mocker): mocker.patch( "agent.nodes.intent_router.get_groq_client" ).return_value.complete_system.return_value = json.dumps( {"intent": "unsupported", "reasoning": "out of scope"} ) from agent.nodes.intent_router import intent_router result = intent_router({**sample_state, "user_query": "what is the weather?"}) assert result["intent"] == "unsupported" def test_defaults_to_sql_on_bad_json(self, sample_state, mocker): mocker.patch( "agent.nodes.intent_router.get_groq_client" ).return_value.complete_system.return_value = "not json at all" from agent.nodes.intent_router import intent_router result = intent_router(sample_state) assert result["intent"] == "sql" # ── SQL Generator ───────────────────────────────────────────────────────────── @pytest.mark.unit class TestSqlGenerator: def test_strips_markdown_fences(self, sample_state, mocker): mocker.patch( "agent.nodes.sql_generator.get_groq_client" ).return_value.complete_system.return_value = ( "```sql\nSELECT * FROM orders LIMIT 5\n```" ) from agent.nodes.sql_generator import sql_generator result = sql_generator(sample_state) assert "```" not in result["generated_code"] assert "SELECT" in result["generated_code"] def test_sets_code_type_to_sql(self, sample_state, mocker): mocker.patch( "agent.nodes.sql_generator.get_groq_client" ).return_value.complete_system.return_value = "SELECT 1" from agent.nodes.sql_generator import sql_generator result = sql_generator(sample_state) assert result["code_type"] == "sql" def test_sets_postgres_dialect_for_neon(self, sample_state, mocker): mocker.patch( "agent.nodes.sql_generator.get_groq_client" ).return_value.complete_system.return_value = "SELECT 1" from agent.nodes.sql_generator import sql_generator result = sql_generator({**sample_state, "connector_id": "neon:public"}) assert result["sql_dialect"] == "postgres" def test_sets_sqlite_dialect_for_csv(self, sample_state, mocker): mocker.patch( "agent.nodes.sql_generator.get_groq_client" ).return_value.complete_system.return_value = "SELECT 1" from agent.nodes.sql_generator import sql_generator result = sql_generator({**sample_state, "connector_id": "csv:http://fake"}) assert result["sql_dialect"] == "sqlite" # ── Pandas Generator ────────────────────────────────────────────────────────── @pytest.mark.unit class TestPandasGenerator: def test_strips_markdown_fences(self, sample_state, mocker): mocker.patch( "agent.nodes.pandas_generator.get_groq_client" ).return_value.complete_system.return_value = ( "```python\nresult = df.head(5)\n```" ) from agent.nodes.pandas_generator import pandas_generator result = pandas_generator(sample_state) assert "```" not in result["generated_code"] assert "result" in result["generated_code"] def test_sets_code_type_to_pandas(self, sample_state, mocker): mocker.patch( "agent.nodes.pandas_generator.get_groq_client" ).return_value.complete_system.return_value = "result = df" from agent.nodes.pandas_generator import pandas_generator result = pandas_generator(sample_state) assert result["code_type"] == "pandas" # ── Safety Validator ────────────────────────────────────────────────────────── @pytest.mark.unit class TestSafetyValidator: def test_passes_valid_sql(self, sample_state): from agent.nodes.safety_validator import safety_validator state = {**sample_state, "generated_code": "SELECT * FROM orders", "code_type": "sql"} result = safety_validator(state) assert result["execution_error"] is None def test_blocks_drop_table(self, sample_state): from agent.nodes.safety_validator import safety_validator state = {**sample_state, "generated_code": "DROP TABLE orders", "code_type": "sql"} result = safety_validator(state) assert result["execution_error"] is not None assert "SAFETY_BLOCK" in result["execution_error"] def test_passes_valid_pandas(self, sample_state): from agent.nodes.safety_validator import safety_validator state = {**sample_state, "generated_code": "result = df.head()", "code_type": "pandas"} result = safety_validator(state) assert result["execution_error"] is None def test_blocks_os_import_in_pandas(self, sample_state): from agent.nodes.safety_validator import safety_validator state = {**sample_state, "generated_code": "import os\nresult = df", "code_type": "pandas"} result = safety_validator(state) assert result["execution_error"] is not None assert "SAFETY_BLOCK" in result["execution_error"] # ── Self Corrector ──────────────────────────────────────────────────────────── @pytest.mark.unit class TestSelfCorrector: def test_increments_correction_attempts(self, sample_state, mocker): mocker.patch( "agent.nodes.self_corrector.get_groq_client" ).return_value.complete_system.return_value = "SELECT 1" from agent.nodes.self_corrector import self_corrector state = {**sample_state, "execution_error": "column not found", "error_class": "nonexistent_column", "correction_attempts": 0} result = self_corrector(state) assert result["correction_attempts"] == 1 def test_clears_execution_error(self, sample_state, mocker): mocker.patch( "agent.nodes.self_corrector.get_groq_client" ).return_value.complete_system.return_value = "SELECT 1" from agent.nodes.self_corrector import self_corrector state = {**sample_state, "execution_error": "syntax error", "error_class": "syntax", "correction_attempts": 1} result = self_corrector(state) assert result["execution_error"] is None def test_gives_up_at_max_attempts(self, sample_state, mocker): mocker.patch("agent.nodes.self_corrector.get_groq_client") from agent.nodes.self_corrector import self_corrector state = {**sample_state, "execution_error": "err", "error_class": "unknown", "correction_attempts": 3, "max_corrections": 3} result = self_corrector(state) assert "unable" in result.get("insight_text", "").lower() def test_strips_markdown_from_corrected_code(self, sample_state, mocker): mocker.patch( "agent.nodes.self_corrector.get_groq_client" ).return_value.complete_system.return_value = "```sql\nSELECT 1\n```" from agent.nodes.self_corrector import self_corrector state = {**sample_state, "execution_error": "err", "error_class": "syntax", "correction_attempts": 0} result = self_corrector(state) assert "```" not in result["generated_code"] # ── Insight Synthesizer ─────────────────────────────────────────────────────── @pytest.mark.unit class TestInsightSynthesizer: def test_returns_insight_text(self, sample_state, sample_rows, mocker): mocker.patch( "agent.nodes.insight_synthesizer.get_groq_client" ).return_value.complete_system.return_value = "Widget leads with $300 in revenue." from agent.nodes.insight_synthesizer import insight_synthesizer state = {**sample_state, "execution_result": sample_rows} result = insight_synthesizer(state) assert result["insight_text"] == "Widget leads with $300 in revenue." def test_no_result_returns_no_results_message(self, sample_state): from agent.nodes.insight_synthesizer import insight_synthesizer state = {**sample_state, "execution_result": None} result = insight_synthesizer(state) assert "no results" in result["insight_text"].lower() def test_empty_result_returns_no_results_message(self, sample_state): from agent.nodes.insight_synthesizer import insight_synthesizer state = {**sample_state, "execution_result": []} result = insight_synthesizer(state) assert "no results" in result["insight_text"].lower()