tests/unit/test_visualizer.py

"""
tests/unit/test_visualizer.py
Tests for chart type inference and Plotly spec generation.
Pure Python — no LLM or DB calls.
"""

import pytest
from agent.nodes.visualizer import visualizer, _infer_chart_type, _build_plotly_spec


# ── Chart type inference ──────────────────────────────────────────────────────

@pytest.mark.unit
class TestInferChartType:
    def test_trend_keyword_gives_line(self):
        rows = [{"month": "Jan", "count": 10}] * 5
        assert _infer_chart_type(rows, "show trend over time") == "line"

    def test_monthly_keyword_gives_line(self):
        rows = [{"month": "Jan", "revenue": 100}] * 12
        assert _infer_chart_type(rows, "monthly revenue breakdown") == "line"

    def test_proportion_keyword_few_rows_gives_pie(self):
        rows = [{"category": x, "pct": i * 10} for i, x in enumerate(["A", "B", "C"])]
        assert _infer_chart_type(rows, "proportion by category") == "pie"

    def test_proportion_many_rows_gives_bar(self):
        rows = [{"region": f"R{i}", "pct": i} for i in range(12)]
        assert _infer_chart_type(rows, "share by region") == "bar"

    def test_scatter_keyword(self):
        rows = [{"x": 1, "y": 2}] * 10
        assert _infer_chart_type(rows, "scatter plot of revenue vs orders") == "scatter"

    def test_two_columns_small_gives_bar(self):
        rows = [{"product": f"P{i}", "total": i * 100} for i in range(5)]
        assert _infer_chart_type(rows, "top products") == "bar"

    def test_many_columns_gives_table(self):
        rows = [{"a": 1, "b": 2, "c": 3, "d": 4}] * 5
        assert _infer_chart_type(rows, "summary") == "table"

    def test_empty_result_gives_table(self):
        assert _infer_chart_type([], "anything") == "table"

    def test_distribution_gives_histogram(self):
        rows = [{"amount": i * 10} for i in range(20)]
        assert _infer_chart_type(rows, "distribution of amounts") == "histogram"


# ── Plotly spec builder ───────────────────────────────────────────────────────

@pytest.mark.unit
class TestBuildPlotlySpec:
    def test_bar_spec_structure(self):
        rows = [{"product": "A", "total": 100}, {"product": "B", "total": 200}]
        spec = _build_plotly_spec(rows, "bar", "top products")
        assert spec["type"] == "bar"
        assert "plotly_json" in spec
        assert spec["plotly_json"]["data"][0]["type"] == "bar"

    def test_line_spec_structure(self):
        rows = [{"month": "Jan", "revenue": 1000}]
        spec = _build_plotly_spec(rows, "line", "monthly revenue")
        assert spec["plotly_json"]["data"][0]["mode"] == "lines+markers"

    def test_pie_spec_structure(self):
        rows = [{"category": "A", "pct": 60}, {"category": "B", "pct": 40}]
        spec = _build_plotly_spec(rows, "pie", "category share")
        assert spec["plotly_json"]["data"][0]["type"] == "pie"

    def test_table_spec_when_type_is_table(self):
        rows = [{"a": 1, "b": 2, "c": 3}]
        spec = _build_plotly_spec(rows, "table", "summary")
        assert spec["type"] == "table"
        assert spec["columns"] == ["a", "b", "c"]
        assert spec["data"] == rows

    def test_empty_result_gives_table(self):
        spec = _build_plotly_spec([], "bar", "test")
        assert spec["type"] == "table"

    def test_layout_has_title(self):
        rows = [{"x": 1, "y": 2}]
        spec = _build_plotly_spec(rows, "bar", "My Query Title")
        assert "My Query Title" in spec["plotly_json"]["layout"]["title"]

    def test_layout_uses_plotly_white_template(self):
        rows = [{"x": 1, "y": 2}]
        spec = _build_plotly_spec(rows, "bar", "test")
        assert spec["plotly_json"]["layout"]["template"] == "plotly_white"


# ── Visualizer node ───────────────────────────────────────────────────────────

@pytest.mark.unit
class TestVisualizerNode:
    def test_adds_chart_spec_to_state(self, sample_state, sample_rows):
        state = {**sample_state, "execution_result": sample_rows}
        result = visualizer(state)
        assert result["chart_spec"] is not None

    def test_no_result_gives_none_chart(self, sample_state):
        state = {**sample_state, "execution_result": None}
        result = visualizer(state)
        assert result["chart_spec"] is None

    def test_empty_result_gives_none_chart(self, sample_state):
        state = {**sample_state, "execution_result": []}
        result = visualizer(state)
        assert result["chart_spec"] is None

    def test_state_keys_preserved(self, sample_state, sample_rows):
        state = {**sample_state, "execution_result": sample_rows}
        result = visualizer(state)
        assert result["session_id"] == sample_state["session_id"]
        assert result["insight_text"] == sample_state["insight_text"]