Spaces:
Running
Running
| import unittest | |
| from datainventor import visualize as V | |
| def sample(n=12): | |
| areas = ["Billing", "Returns", "Shipping"] | |
| return [{"id": i, "prompt": f"Question number {i} about an order that arrived late and damaged", | |
| "completion": " ".join(["word"] * (20 + 10 * i)), "area": areas[i % 3], "difficulty": "basic" if i % 2 else "advanced", | |
| "score": str(i % 5 + 1), "ok": bool(i % 2)} for i in range(n)] | |
| class Prepare(unittest.TestCase): | |
| def test_types_and_derived_lengths(self): | |
| rows, cols = V.prepare(sample()) | |
| self.assertEqual(cols["prompt"]["type"], "text") | |
| self.assertEqual(cols["area"]["type"], "cat") | |
| self.assertEqual(cols["score"]["type"], "num") | |
| self.assertEqual(cols["ok"]["type"], "cat") | |
| self.assertNotIn("id", cols) | |
| self.assertEqual(cols["completion words"]["derived"], "words") | |
| self.assertEqual(rows[0]["_id"], 1) | |
| self.assertEqual(rows[0]["completion words"], 20.0) | |
| def test_chat_cells_become_text_with_turns(self): | |
| rows, cols = V.prepare([{"messages": [{"role": "user", "content": "hi " * 30}, {"role": "assistant", "content": "hello " * 30}]}] * 3) | |
| self.assertEqual(cols["messages"]["type"], "text") | |
| self.assertEqual(rows[0]["messages turns"], 2.0) | |
| class Charts(unittest.TestCase): | |
| def setUp(self): | |
| self.rows, self.cols = V.prepare(sample(30)) | |
| def test_every_row_lands_in_exactly_one_bin(self): | |
| c = V.chart(self.rows, self.cols, {"type": "hist", "x": "completion words"}) | |
| ids = [i for b in c["marks"] for i in b["ids"]] | |
| self.assertEqual(sorted(ids), list(range(1, 31))) | |
| def test_bar_counts_and_ids(self): | |
| c = V.chart(self.rows, self.cols, {"type": "bar", "x": "area", "agg": "count"}) | |
| self.assertEqual(sum(m["value"] for m in c["marks"]), 30) | |
| self.assertTrue(all(len(m["ids"]) == m["value"] for m in c["marks"])) | |
| def test_box_quartiles_are_ordered(self): | |
| c = V.chart(self.rows, self.cols, {"type": "box", "x": "area", "y": "completion words"}) | |
| for m in c["marks"]: | |
| self.assertTrue(m["lo"] <= m["q1"] <= m["median"] <= m["q3"] <= m["hi"]) | |
| def test_check_spec_fixes_and_rejects(self): | |
| self.assertEqual(V.check_spec({"type": "bar", "x": "completion words"}, self.cols)["type"], "hist") | |
| self.assertEqual(V.check_spec({"type": "bar", "x": "area", "y": "score"}, self.cols)["agg"], "mean") | |
| with self.assertRaises(V.PlanError): | |
| V.check_spec({"type": "box", "x": "area", "y": "area"}, self.cols) | |
| with self.assertRaises(V.PlanError): | |
| V.check_spec({"type": "bar", "x": "missing"}, self.cols) | |
| class Plans(unittest.TestCase): | |
| def setUp(self): | |
| self.rows, self.cols = V.prepare(sample()) | |
| def test_parse_plan_with_fences_and_derived_column(self): | |
| p = V.parse_plan('Sure!\n```json\n{"title": "Tone by area", "derive": [{"name": "Tone!", "kind": "label", "from": "completion", "labels": ["calm", "curt"]}],' | |
| ' "chart": {"type": "stack", "x": "area", "color": "tone"}, "follow_ups": ["a"]}\n```', self.cols) | |
| self.assertEqual(p["derive"][0]["name"], "tone") | |
| self.assertEqual(p["chart"]["agg"], "share") | |
| def test_derived_name_never_overwrites_an_original_column(self): | |
| p = V.parse_plan('{"title": "t", "derive": [{"name": "area", "kind": "label", "from": "prompt", "labels": ["a", "b"]}],' | |
| ' "chart": {"type": "bar", "x": "area"}}', self.cols) | |
| self.assertEqual(p["derive"][0]["name"], "area (model)") | |
| self.assertEqual(p["chart"]["x"], "area (model)") | |
| def test_snake_case_names_still_match_the_chart(self): | |
| p = V.parse_plan('{"derive": [{"name": "step_by_step", "kind": "pattern", "from": "completion", "pattern": "(?i)step \\\\d"}],' | |
| ' "chart": {"type": "bar", "x": "step_by_step"}}', self.cols) | |
| self.assertEqual(p["derive"][0]["name"], "step by step") | |
| self.assertEqual(p["chart"]["x"], "step by step") | |
| self.assertEqual(p["derive"][0]["pattern"], "step \\d") | |
| def test_dangerous_pattern_rejected(self): | |
| with self.assertRaises(V.PlanError): | |
| V.parse_plan('{"derive": [{"name": "x", "kind": "pattern", "from": "prompt", "pattern": "(a+)+$"}], "chart": {"type": "bar", "x": "x"}}', self.cols) | |
| def test_label_reading(self): | |
| d = {"kind": "label", "labels": ["Formal", "Friendly"]} | |
| self.assertEqual(V._read_labels('{"1": "formal", "2": "friendly tone"}', d, 2), {1: "Formal", 2: "Friendly"}) | |
| self.assertIsNone(V._read_labels('{"1": "angry"}', d, 2)) | |
| def test_pattern_fill_needs_no_model(self): | |
| rows, _ = V.prepare([{"t": "Step 1 do this. Step 2 do that."}, {"t": "Just a plain answer here."}]) | |
| vals = V.fill({"name": "steps", "kind": "pattern", "from": "t", "pattern": r"step \d"}, rows, chat=None) | |
| self.assertEqual(vals, {1: "yes", 2: "no"}) | |
| if __name__ == "__main__": | |
| unittest.main() | |