Spaces:
Running
Running
File size: 5,049 Bytes
917845f | 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 | 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()
|