From c73aec959d6cfcba8383d7a71161edbcf6770838 Mon Sep 17 00:00:00 2001
From: mhostetter <matthostetter@gmail.com>
Date: Tue, 8 Feb 2022 12:20:40 -0500
Subject: [PATCH] Add more `is_prime_power()` unit tests

---
 scripts/generate_int_test_vectors.py |  22 +++++++++++++++++++++-
 tests/conftest.py                    |   5 +++++
 tests/data/is_prime_power.pkl        | Bin 0 -> 239 bytes
 tests/test_factor.py                 |   9 ---------
 tests/test_primes.py                 |  11 +++++++++++
 5 files changed, 37 insertions(+), 10 deletions(-)
 create mode 100644 tests/data/is_prime_power.pkl

diff --git a/scripts/generate_int_test_vectors.py b/scripts/generate_int_test_vectors.py
index 1b60673d8..3488c17ca 100644
--- a/scripts/generate_int_test_vectors.py
+++ b/scripts/generate_int_test_vectors.py
@@ -11,7 +11,7 @@
 
 import sage
 import numpy as np
-from sage.all import Integer, Integers, xgcd, lcm, prod, isqrt, log, crt, euler_phi, prime_range, nth_prime, previous_prime, next_prime, is_prime
+from sage.all import Integer, Integers, xgcd, lcm, prod, isqrt, log, crt, euler_phi, prime_range, nth_prime, previous_prime, next_prime, is_prime, is_prime_power
 from sage.crypto.util import carmichael_lambda
 
 PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tests")
@@ -230,3 +230,23 @@ def save_pickle(d, folder, name):
     Z[i] = bool(z)
 d = {"X": X, "Z": Z}
 save_pickle(d, FOLDER, "is_prime.pkl")
+
+# set_seed(SEED + 306)
+# X = [random.randint(-100, 100) for _ in range(20)] + [random.randint(100, 1_000_000_000) for _ in range(20)]
+# Z = [0,]*len(X)
+# for i in range(len(X)):
+#     x = X[i]
+#     z = is_prime(x)
+#     Z[i] = bool(z)
+# d = {"X": X, "Z": Z}
+# save_pickle(d, FOLDER, "is_composite.pkl")
+
+set_seed(SEED + 307)
+X = [random.randint(-100, 100) for _ in range(20)] + [random.randint(100, 1_000_000_000) for _ in range(20)]
+Z = [0,]*len(X)
+for i in range(len(X)):
+    x = X[i]
+    z = is_prime_power(x)
+    Z[i] = bool(z)
+d = {"X": X, "Z": Z}
+save_pickle(d, FOLDER, "is_prime_power.pkl")
diff --git a/tests/conftest.py b/tests/conftest.py
index 6dd1f574f..b9ad3f89f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -111,3 +111,8 @@ def next_prime():
 @pytest.fixture(scope="session")
 def is_prime():
     return read_pickle("is_prime.pkl")
+
+
+@pytest.fixture(scope="session")
+def is_prime_power():
+    return read_pickle("is_prime_power.pkl")
diff --git a/tests/data/is_prime_power.pkl b/tests/data/is_prime_power.pkl
new file mode 100644
index 0000000000000000000000000000000000000000..02e6aca07eb4335c764948b90a4e9b7c0703e14b
GIT binary patch
literal 239
zcmZo*nfinQ0&1sd^e{$DiJhY19q)Db|NsBqir&oL4BkrK3f=)gp*cVyuNxp@K8V;3
zB)oaO=7K13uXiBI&1)G*W)+BFm7JmB_2iYQh?o94Rt_(p<`4Q_>vP|6c%4qar{{HQ
z?|d__4<|QsdL52o;_~|X=dQWehei3!UfFHUYF^v>ZH>H&PAeLD^*xDX_L{z`!qn?0
uixaz7&<TAmuc>xZSiR<E&0|dMVT=Mfx3jae1B`$SAPdOG1R%lGQau1sFmZJN

literal 0
HcmV?d00001

diff --git a/tests/test_factor.py b/tests/test_factor.py
index dec74828a..17de7ced9 100644
--- a/tests/test_factor.py
+++ b/tests/test_factor.py
@@ -246,15 +246,6 @@ def test_divisor_sigma():
     assert [galois.divisor_sigma(n, k=3) for n in range(1, 43)] == sigma_3
 
 
-def test_is_prime_power():
-    # https://oeis.org/A246655
-    prime_powers = np.array([2,3,4,5,7,8,9,11,13,16,17,19,23,25,27,29,31,32,37,41,43,47,49,53,59,61,64,67,71,73,79,81,83,89,97,101,103,107,109,113,121,125,127,128,131,137,139,149,151,157,163,167,169,173,179,181,191,193,197,199,211])
-    n = np.arange(1, prime_powers[-1] + 1)
-    is_prime_power = np.zeros(n.size, dtype=bool)
-    is_prime_power[prime_powers - 1] = True  # -1 for 1-indexed
-    assert [galois.is_prime_power(ni) for ni in n] == is_prime_power.tolist()
-
-
 def test_is_perfect_power():
     # https://oeis.org/A001597
     perfect_powers = np.array([1,4,8,9,16,25,27,32,36,49,64,81,100,121,125,128,144,169,196,216,225,243,256,289,324,343,361,400,441,484,512,529,576,625,676,729,784,841,900,961,1000,1024,1089,1156,1225,1296,1331,1369,1444,1521,1600,1681,1728,1764])
diff --git a/tests/test_primes.py b/tests/test_primes.py
index 794cea854..751796499 100644
--- a/tests/test_primes.py
+++ b/tests/test_primes.py
@@ -143,6 +143,17 @@ def test_is_composite():
     assert [galois.is_composite(ni) for ni in n] == is_composite.tolist()
 
 
+def test_is_prime_power_exceptions():
+    with pytest.raises(TypeError):
+        galois.is_prime_power(13.0)
+
+
+def test_is_prime_power(is_prime_power):
+    X, Z = is_prime_power["X"], is_prime_power["Z"]
+    for i in range(len(X)):
+        assert galois.is_prime_power(X[i]) == Z[i]
+
+
 def test_fermat_primality_test_exceptions():
     with pytest.raises(TypeError):
         galois.fermat_primality_test(13.0)