tally

commit acaf828620f7369a99cdddc3efd4709e6725aff3
parent bfca0968cceb76b8d1602c13f8627017de3efd9f
author Owen Rusk <owen@papermothgames.com>
date   2024-06-13 13:40:08 -0500

tests for split and combine

diff --git a/tests/test_shamir.py b/tests/test_shamir.py
new file mode 100644
index 0000000..c3b9a42
--- /dev/null
+++ b/tests/test_shamir.py
@@ -0,0 +1,38 @@
+import itertools
+import os
+import unittest
+
+from tally import shamir
+
+
+class ShamirTest(unittest.TestCase):
+    def test_round_trip(self) -> None:
+        secret = b"not a real secret"
+        shares = shamir.split(secret, 3, 5)
+        self.assertEqual(len(shares), 5)
+        self.assertEqual(shamir.combine(shares[:3]), secret)
+
+    def test_any_three_of_five(self) -> None:
+        secret = os.urandom(32)
+        shares = shamir.split(secret, 3, 5)
+        for subset in itertools.combinations(shares, 3):
+            self.assertEqual(shamir.combine(list(subset)), secret)
+
+    def test_two_of_five_is_not_enough(self) -> None:
+        secret = os.urandom(32)
+        shares = shamir.split(secret, 3, 5)
+        self.assertNotEqual(shamir.combine(shares[:2]), secret)
+
+    def test_bad_k_and_n(self) -> None:
+        for k, n in [(1, 3), (4, 3), (2, 256), (0, 0)]:
+            with self.assertRaises(ValueError):
+                shamir.split(b"x", k, n)
+
+    def test_same_x_twice(self) -> None:
+        shares = shamir.split(b"x", 2, 3)
+        with self.assertRaises(ValueError):
+            shamir.combine([shares[0], shares[0]])
+
+
+if __name__ == "__main__":
+    unittest.main()