beam: convert test_levenshtein to unittest.TestCase so lane gate collects it

wizard-work/beam/test_levenshtein.py

@@ -1,63 +1,71 @@
-import random, string
+import random
+import unittest
+
 from levenshtein import edit_distance, ops
 
 
 def brute(a, b):
     n, m = len(a), len(b)
-    dp = [[0]*(m+1) for _ in range(n+1)]
-    for i in range(n+1): dp[i][0]=i
-    for j in range(m+1): dp[0][j]=j
-    for i in range(1,n+1):
-        for j in range(1,m+1):
-            c = 0 if a[i-1]==b[j-1] else 1
-            dp[i][j]=min(dp[i-1][j]+1, dp[i][j-1]+1, dp[i-1][j-1]+c)
+    dp = [[0] * (m + 1) for _ in range(n + 1)]
+    for i in range(n + 1):
+        dp[i][0] = i
+    for j in range(m + 1):
+        dp[0][j] = j
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            c = 0 if a[i - 1] == b[j - 1] else 1
+            dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j - 1] + 1, dp[i - 1][j - 1] + c)
     return dp[n][m]
 
 
-def test_known():
-    cases = [("kitten","sitting",3),("","abc",3),("abc","",3),
-             ("flaw","lawn",2),("same","same",0),("a","b",1),
-             ("","",-0 if False else 0)]
-    for a,b,e in cases:
-        assert edit_distance(a,b)==e, (a,b,edit_distance(a,b),e)
+class TestLevenshtein(unittest.TestCase):
+    def test_known(self):
+        cases = [("kitten", "sitting", 3), ("", "abc", 3), ("abc", "", 3),
+                 ("flaw", "lawn", 2), ("same", "same", 0), ("a", "b", 1),
+                 ("", "", 0)]
+        for a, b, e in cases:
+            with self.subTest(a=a, b=b):
+                self.assertEqual(edit_distance(a, b), e)
 
+    def test_random_vs_brute(self):
+        rng = random.Random(42)
+        alpha = "ab"
+        for _ in range(2000):
+            la = rng.randint(0, 8)
+            lb = rng.randint(0, 8)
+            a = "".join(rng.choice(alpha) for _ in range(la))
+            b = "".join(rng.choice(alpha) for _ in range(lb))
+            with self.subTest(a=a, b=b):
+                self.assertEqual(edit_distance(a, b), brute(a, b))
 
-def test_random_vs_brute():
-    rng = random.Random(42)
-    alpha = "ab"
-    for _ in range(2000):
-        la = rng.randint(0,8); lb = rng.randint(0,8)
-        a = "".join(rng.choice(alpha) for _ in range(la))
-        b = "".join(rng.choice(alpha) for _ in range(lb))
-        assert edit_distance(a,b)==brute(a,b), (a,b)
-
-
-def test_ops_roundtrip():
-    rng = random.Random(7)
-    for _ in range(500):
-        a = "".join(rng.choice("ab") for _ in range(rng.randint(0,6)))
-        b = "".join(rng.choice("ab") for _ in range(rng.randint(0,6)))
-        o = ops(a,b)
-        # apply ops to a -> should yield b
-        out = []
-        i = 0  # index into a
-        j = 0  # index into b
-        for tag in o:
-            if tag == "=":
-                assert a[i] == b[j]
-                out.append(a[i]); i += 1; j += 1
-            elif tag == "S":
-                out.append(b[j]); i += 1; j += 1
-            elif tag == "D":
-                i += 1
-            elif tag == "I":
-                # need current b char; find next unconsumed b
-                out.append(b[j]); j += 1
-        assert "".join(out) == b, (a,b,o,"".join(out))
-        edits = sum(1 for t in o if t in ("S","D","I"))
-        assert edits == edit_distance(a,b), (a,b,edits,edit_distance(a,b))
+    def test_ops_roundtrip(self):
+        rng = random.Random(7)
+        for _ in range(500):
+            a = "".join(rng.choice("ab") for _ in range(rng.randint(0, 6)))
+            b = "".join(rng.choice("ab") for _ in range(rng.randint(0, 6)))
+            o = ops(a, b)
+            out = []
+            i = 0
+            j = 0
+            for tag in o:
+                if tag == "=":
+                    self.assertEqual(a[i], b[j])
+                    out.append(a[i])
+                    i += 1
+                    j += 1
+                elif tag == "S":
+                    out.append(b[j])
+                    i += 1
+                    j += 1
+                elif tag == "D":
+                    i += 1
+                elif tag == "I":
+                    out.append(b[j])
+                    j += 1
+            self.assertEqual("".join(out), b, (a, b, o, "".join(out)))
+            edits = sum(1 for t in o if t in ("S", "D", "I"))
+            self.assertEqual(edits, edit_distance(a, b), (a, b, edits, edit_distance(a, b)))
 
 
 if __name__ == "__main__":
-    test_known(); test_random_vs_brute(); test_ops_roundtrip()
-    print("ALL OK")
+    unittest.main()