test_gradio.py

#!/usr/bin/env python3
"""
Standalone tests for the Gradio interface.
Runs without launching the server or needing Ollama/Neon.
"""

import os
import sys
import unittest
from pathlib import Path
from unittest.mock import patch

# Setup path and env before imports
os.environ.setdefault("MPLCONFIGDIR", str(Path(__file__).parent / ".matplotlib"))
sys.path.insert(0, str(Path(__file__).parent))

import pandas as pd

from agent_runner import format_result_for_chat, safe_df


class TestAgentRunner(unittest.TestCase):
    """Tests for agent_runner helpers used by Gradio."""

    def test_format_result_empty(self):
        text, df, viz = format_result_for_chat({})
        self.assertIn("No response data", text)
        self.assertIsNone(df)
        self.assertIsNone(viz)

    def test_format_result_non_dict(self):
        text, df, viz = format_result_for_chat("plain string")
        self.assertEqual(text, "plain string")
        self.assertIsNone(df)
        self.assertIsNone(viz)

    def test_format_result_with_final_answer_and_table(self):
        result = {
            "final_answer": "Here are the results.",
            "query_result": [{"col_a": 1, "col_b": 2}, {"col_a": 3, "col_b": 4}],
        }
        text, df, viz = format_result_for_chat(result)
        self.assertIn("Here are the results", text)
        self.assertIn("**Results**", text)
        self.assertIn("| col_a | col_b |", text)
        self.assertIsInstance(df, pd.DataFrame)
        self.assertEqual(len(df), 2)
        self.assertIsNone(viz)

    def test_format_result_with_sql_and_error(self):
        result = {
            "datasource": "sql",
            "sql_query": "SELECT * FROM t",
            "query_error": "table not found",
        }
        text, df, viz = format_result_for_chat(result)
        self.assertIn("Agent Reasoning", text)
        self.assertIn("SELECT * FROM t", text)
        self.assertIn("table not found", text)

    def test_safe_df(self):
        out = safe_df([{"a": 1}, {"a": 2}])
        self.assertIsInstance(out, pd.DataFrame)
        self.assertEqual(list(out["a"]), [1, 2])


class TestGradioUI(unittest.TestCase):
    """Tests for Gradio UI handlers."""

    def setUp(self):
        # Import here to avoid loading Gradio/UI until needed
        from app_gradio import chat_turn, clear_chat, approve_click, reject_click

        self.chat_turn = chat_turn
        self.clear_chat = clear_chat
        self.approve_click = approve_click
        self.reject_click = reject_click

    def test_chat_turn_no_db_url(self):
        """When db_url is empty and no env fallback, returns config message without calling agent."""
        with patch.dict(os.environ, {"NEON_DATABASE_URL": "", "DATABASE_URL": ""}, clear=False):
            hist, tid, awaiting, pending, url, df, viz = self.chat_turn(
                "What are sales?",
                [],
                "thread-123",
                False,
                None,
                "",  # empty db_url
                "llama3.2",
            )
        self.assertEqual(len(hist), 1)
        self.assertIn("Set Neon connection string", hist[0][1])
        self.assertFalse(awaiting)
        self.assertIsNone(pending)
        self.assertIsNone(df)
        self.assertIsNone(viz)

    def test_clear_chat(self):
        """clear_chat returns fresh state and empty table/chart."""
        out = self.clear_chat("old-thread")
        self.assertEqual(len(out), 8)
        hist, tid, awaiting, pending, msg, approval_upd, table, chart = out
        self.assertEqual(hist, [])
        self.assertNotEqual(tid, "old-thread")
        self.assertFalse(awaiting)
        self.assertIsNone(pending)
        self.assertEqual(msg, "")
        self.assertTrue(table.empty)
        self.assertIsNone(chart)

    def test_approve_click_no_pending(self):
        """approve_click with no pending result returns early."""
        hist, awaiting, pending, df, viz = self.approve_click(
            [("u", "a")], None, "tid", "postgres://...", "llama3.2"
        )
        self.assertEqual(len(hist), 1)
        self.assertFalse(awaiting)
        self.assertIsNone(df)
        self.assertIsNone(viz)

    def test_reject_click_no_pending(self):
        """reject_click with no pending result returns early."""
        hist, awaiting, pending, df, viz = self.reject_click(
            [("u", "a")], None, "tid", "postgres://...", "llama3.2"
        )
        self.assertEqual(len(hist), 1)
        self.assertFalse(awaiting)
        self.assertIsNone(df)
        self.assertIsNone(viz)

    def test_chat_turn_returns_sql_output_and_chart(self):
        """When agent returns query_result and chart, chat_turn must return df and viz. If not, tests fail."""
        mock_result = {
            "final_answer": "Here are sales by category.",
            "query_result": [{"category": "A", "total": 100}, {"category": "B", "total": 200}],
            "visualization_path": "/tmp/test_chart.png",
        }
        with patch("app_gradio.run_agent", return_value=(mock_result, False)):
            with patch("agent_runner.Path") as mock_path_cls:
                mock_path_cls.return_value.exists.return_value = True
                hist, tid, awaiting, pending, url, df, viz_path = self.chat_turn(
                    "Sales by category", [], "tid", False, None, "postgres://x", "llama3.2"
                )
        self.assertFalse(df is None or df.empty, "SQL output must show up in table - df must not be empty")
        self.assertIsNotNone(viz_path, "Chart must show up - viz_path must not be None")

    def test_approve_click_returns_sql_output_and_chart(self):
        """When approve runs and agent returns query_result and chart, approve_click must return df and viz."""
        mock_result = {
            "final_answer": "Approved and executed.",
            "query_result": [{"region": "North", "sales": 500}],
            "visualization_path": "/tmp/approved_chart.png",
        }
        with patch("app_gradio.resume_agent", return_value=mock_result):
            with patch("agent_runner.Path") as mock_path_cls:
                mock_path_cls.return_value.exists.return_value = True
                hist, awaiting, pending, df, viz_path = self.approve_click(
                    [("u", "a")], {"sql_query": "SELECT 1"}, "tid", "postgres://x", "llama3.2"
                )
        self.assertFalse(df is None or df.empty, "SQL output must show up after approve - df must not be empty")
        self.assertIsNotNone(viz_path, "Chart must show up after approve - viz_path must not be None")


class TestGradioBlocks(unittest.TestCase):
    """Tests that the Gradio UI builds without error."""

    def test_build_ui(self):
        """build_ui returns a Blocks instance."""
        from app_gradio import build_ui
        import gradio as gr

        demo = build_ui()
        self.assertIsInstance(demo, gr.Blocks)

    def test_sample_questions_present(self):
        """Sample questions list is non-empty."""
        from app_gradio import SAMPLE_QUESTIONS

        self.assertGreater(len(SAMPLE_QUESTIONS), 0)
        self.assertIn("Top 5 products by total sales", SAMPLE_QUESTIONS)
        self.assertIn("How many PTO days do we get?", SAMPLE_QUESTIONS)


if __name__ == "__main__":
    unittest.main(verbosity=2)