| """ |
| 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 |
|
|
|
|
| |
|
|
| @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" |
|
|
|
|
| |
|
|
| @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" |
|
|
|
|
| |
|
|
| @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"] |
|
|
