test_numbertheory: test square_root_mod_prime

extend the testing for large primes with hypothesis

use pytest.parametrize to test with small primes

Author: tomato42

src/ecdsa/test_numbertheory.py

@@ -4,6 +4,7 @@
 except ImportError:
     import unittest
 import hypothesis.strategies as st
+import pytest
 from hypothesis import given, settings, example
 try:
     from hypothesis import HealthCheck
@@ -56,40 +57,22 @@ def test_numbertheory():
   for i in range(len(bigprimes) - 1):
     assert next_prime(bigprimes[i]) == bigprimes[i + 1]
 
-  error_tally = 0
-
-  # Test the square_root_mod_prime function:
-
-  for p in smallprimes:
-    print_("Testing square_root_mod_prime for modulus p = %d." % p)
-    squares = []
-
-    for root in range(0, 1 + p // 2):
-      sq = (root * root) % p
-      squares.append(sq)
-      calculated = square_root_mod_prime(sq, p)
-      if (calculated * calculated) % p != sq:
-        error_tally = error_tally + 1
-        print_("Failed to find %d as sqrt( %d ) mod %d. Said %d." % \
-               (root, sq, p, calculated))
-
-    for nonsquare in range(0, p):
-      if nonsquare not in squares:
-        try:
-          calculated = square_root_mod_prime(nonsquare, p)
-        except SquareRootError:
-          pass
-        else:
-          error_tally = error_tally + 1
-          print_("Failed to report no root for sqrt( %d ) mod %d." % \
-                 (nonsquare, p))
 
-  class FailedTest(Exception):
-    pass
+@pytest.mark.parametrize("prime", smallprimes)
+def test_square_root_mod_prime_for_small_primes(prime):
+    squares = set()
+    for num in range(0, 1 + prime // 2):
+        sq = num * num % prime
+        squares.add(sq)
+        root = square_root_mod_prime(sq, prime)
+        # tested for real with TestNumbertheory.test_square_root_mod_prime
+        assert root * root % prime == sq
 
-  print_(error_tally, "errors detected.")
-  if error_tally != 0:
-    raise FailedTest("%d errors detected" % error_tally)
+    for nonsquare in range(0, prime):
+        if nonsquare in squares:
+            continue
+        with pytest.raises(SquareRootError):
+            square_root_mod_prime(nonsquare, prime)
 
 
 @st.composite
@@ -102,6 +85,24 @@ def st_two_nums_rel_prime(draw):
     return num, mod
 
 
+@st.composite
+def st_primes(draw, *args, **kwargs):
+    if "min_value" not in kwargs:
+        kwargs["min_value"] = 1
+    prime = draw(st.sampled_from(smallprimes) |
+                 st.integers(*args, **kwargs)
+                 .filter(lambda x: is_prime(x)))
+    return prime
+
+
+@st.composite
+def st_num_square_prime(draw):
+    prime = draw(st_primes(max_value=2**1024))
+    num = draw(st.integers(min_value=0, max_value=1 + prime // 2))
+    sq = num * num % prime
+    return sq, prime
+
+
 HYP_SETTINGS = {}
 if HC_PRESENT:
     HYP_SETTINGS['suppress_health_check']=[HealthCheck.filter_too_much,
@@ -111,6 +112,17 @@ def st_two_nums_rel_prime(draw):
 
 
 class TestNumbertheory(unittest.TestCase):
+    @unittest.skipUnless(HC_PRESENT,
+                         "Hypothesis 2.0.0 can't be made tolerant of hard to "
+                         "meet requirements (like `is_prime()`)")
+    @settings(**HYP_SETTINGS)
+    @given(st_num_square_prime())
+    def test_square_root_mod_prime(self, vals):
+        square, prime = vals
+
+        calc = square_root_mod_prime(square, prime)
+        assert calc * calc % prime == square
+
     @settings(**HYP_SETTINGS)
     @given(st.integers(min_value=1, max_value=10**12))
     @example(265399 * 1526929)