Add DigestFunctionTest unittest

pyauth · Aug 31, 2024 · 516531a · 516531a
1 parent 2ed7bec
commit 516531a
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/test.py b/test.py
@@ -334,6 +334,16 @@ def test_valid_window(self):
         self.assertTrue(totp.verify("681610", 200, 1))
         self.assertFalse(totp.verify("195979", 200, 1))
 
+class DigestFunctionTest(unittest.TestCase):
+    def test_md5(self):
+        with self.assertRaises(ValueError) as cm:
+            pyotp.OTP(s="secret", digest=hashlib.md5)
+        self.assertEqual("selected digest function must generate digest size greater than or equals to 18 bytes", str(cm.exception))
+
+    def test_shake128(self):
+        with self.assertRaises(ValueError) as cm:
+            pyotp.OTP(s="secret", digest=hashlib.shake_128)
+        self.assertEqual("selected digest function must generate digest size greater than or equals to 18 bytes", str(cm.exception))
 
 class ParseUriTest(unittest.TestCase):
     def test_invalids(self):