File size: 5,095 Bytes
abd4352 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
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"]
|