study-buddy / tests /test_composition_compiler.py
GitHub Actions
deploy d092bea3608b7a29952f16357fda39b7a29e399b
2e818da
Raw
History Blame Contribute Delete
4.26 kB
from app.schemas.visual_lesson import ExtractedChartDataset, ExtractedChartRow
from app.services.composition_compiler import CompositionCompiler
def test_chart_without_extracted_data_falls_back_to_illustrative_values():
compiler = CompositionCompiler()
payload = compiler.build("project-1", "give me a bar chart")
dataset = payload.spec.datasets[0]
assert "illustrative" in payload.spec.answer_markdown.lower()
assert payload.spec.assumptions == ["The chart dataset is illustrative."]
assert len(dataset.rows) > 0
def test_chart_uses_grounded_rows_with_categorical_labels():
compiler = CompositionCompiler()
extracted = ExtractedChartDataset(
has_data=True,
chart_family="bar",
x_label="Method",
y_label="Multi-subject average",
rows=[
ExtractedChartRow(label="REINFORCE", value=22.6),
ExtractedChartRow(label="RLOO (ours)", value=31.2),
],
)
payload = compiler.build("project-1", "I need a plot for visualizing performance difference", extracted=extracted)
dataset = payload.spec.datasets[0]
chart = payload.spec.panels[0].chart
assert [row["x"] for row in dataset.rows] == ["REINFORCE", "RLOO (ours)"]
assert [row["y"] for row in dataset.rows] == [22.6, 31.2]
assert chart.x_label == "Method"
assert chart.y_label == "Multi-subject average"
x_encoding = next(e for e in chart.encodings if e.channel == "x")
assert x_encoding.value_type == "nominal"
assert payload.spec.assumptions == []
assert "transcribed from your selected context" in payload.spec.answer_markdown
assert "grounded in selected context" in payload.spec.provenance.interpretation
def test_chart_uses_grounded_rows_with_numeric_labels_as_quantitative():
compiler = CompositionCompiler()
extracted = ExtractedChartDataset(
has_data=True,
chart_family="line",
rows=[ExtractedChartRow(label="2020", value=10.0), ExtractedChartRow(label="2021", value=15.0)],
)
payload = compiler.build("project-1", "line chart of results", extracted=extracted)
chart = payload.spec.panels[0].chart
assert chart.family == "line"
x_encoding = next(e for e in chart.encodings if e.channel == "x")
assert x_encoding.value_type == "quantitative"
dataset = payload.spec.datasets[0]
assert dataset.rows[0]["x"] == 2020.0
def test_grounded_categorical_data_does_not_default_to_line_when_prompt_has_no_chart_keyword():
"""Regression test: a prompt like 'I need a table for this plot' or 'plot form for easy
visualization' contains no bar/line/scatter keyword, so the old keyword-only heuristic always
fell through to 'line' even for a table of unordered named categories."""
compiler = CompositionCompiler()
extracted = ExtractedChartDataset(
has_data=True,
chart_family="bar",
rows=[
ExtractedChartRow(label="Basic Medicine", value=9.9),
ExtractedChartRow(label="Law", value=9.2),
ExtractedChartRow(label="Economics", value=6.7),
],
)
payload = compiler.build("project-1", "i need a table for this plot", extracted=extracted)
assert payload.spec.panels[0].chart.family == "bar"
def test_explicit_force_graph_keyword_overrides_extracted_chart_family():
compiler = CompositionCompiler()
extracted = ExtractedChartDataset(has_data=True, chart_family="bar", rows=[ExtractedChartRow(label="A", value=1.0)])
payload = compiler.build("project-1", "force graph of the network", extracted=extracted)
assert payload.spec.panels[0].chart.family == "force-graph"
assert payload.spec.datasets[0].columns == ["source", "target", "weight"]
def test_build_routes_to_chart_when_extracted_data_present_without_chart_keyword():
compiler = CompositionCompiler()
extracted = ExtractedChartDataset(
has_data=True, rows=[ExtractedChartRow(label="A", value=1.0), ExtractedChartRow(label="B", value=2.0)]
)
payload = compiler.build("project-1", "visualize the performance difference between methods", extracted=extracted)
assert len(payload.spec.panels) == 1
assert payload.spec.panels[0].chart.family in {"line", "bar", "scatter", "heatmap"}