Merge pull request #88 from GiacomoPope/include_more_tests

include more tests for decaps in mlkem
GiacomoPope · Nov 26, 2024 · 038a5d9 · 038a5d9
2 parents a0dac69 + b9a3ced
commit 038a5d9
Showing 1 changed file with 34 additions and 7 deletions.
diff --git a/tests/test_ml_kem.py b/tests/test_ml_kem.py
@@ -105,7 +105,7 @@ def generic_keygen_kat(self, ML_KEM, index):
             self.assertEqual(ek, ek_kat)
             self.assertEqual(dk, dk_kat)
 
-    def generic_encap_decap_kat(self, ML_KEM, index):
+    def generic_encap_kat(self, ML_KEM, index):
         with open(
             "assets/ML-KEM-encapDecap-FIPS203/internalProjection.json"
         ) as f:
@@ -126,6 +126,24 @@ def generic_encap_decap_kat(self, ML_KEM, index):
             K_prime = ML_KEM.decaps(dk_kat, c_kat)
             self.assertEqual(K_prime, k_kat)
 
+    def generic_decap_kat(self, ML_KEM, index):
+        with open(
+            "assets/ML-KEM-encapDecap-FIPS203/internalProjection.json"
+        ) as f:
+            data = json.load(f)
+        kat_data = data["testGroups"][3 + index]["tests"]
+
+        # Parse out the decaps key
+        dk_hex = data["testGroups"][3 + index]["dk"]
+        dk_kat = bytes.fromhex(dk_hex)
+
+        # Ensure that decaps works
+        for test in kat_data:
+            c_kat = bytes.fromhex(test["c"])
+            k_kat = bytes.fromhex(test["k"])
+            K = ML_KEM.decaps(dk_kat, c_kat)
+            self.assertEqual(K, k_kat)
+
     def test_ML_KEM_512_keygen(self):
         self.generic_keygen_kat(ML_KEM_512, 0)
 
@@ -135,11 +153,20 @@ def test_ML_KEM_768_keygen(self):
     def test_ML_KEM_1024_keygen(self):
         self.generic_keygen_kat(ML_KEM_1024, 2)
 
-    def test_ML_KEM_512_encap_decap(self):
-        self.generic_encap_decap_kat(ML_KEM_512, 0)
+    def test_ML_KEM_512_encap(self):
+        self.generic_encap_kat(ML_KEM_512, 0)
+
+    def test_ML_KEM_768_encap(self):
+        self.generic_encap_kat(ML_KEM_768, 1)
+
+    def test_ML_KEM_1024_encap(self):
+        self.generic_encap_kat(ML_KEM_1024, 2)
+
+    def test_ML_KEM_512_decap(self):
+        self.generic_decap_kat(ML_KEM_512, 0)
 
-    def test_ML_KEM_768_encap_decap(self):
-        self.generic_encap_decap_kat(ML_KEM_768, 1)
+    def test_ML_KEM_768_decap(self):
+        self.generic_decap_kat(ML_KEM_768, 1)
 
-    def test_ML_KEM_1024_encap_decap(self):
-        self.generic_encap_decap_kat(ML_KEM_1024, 2)
+    def test_ML_KEM_1024_decap(self):
+        self.generic_decap_kat(ML_KEM_1024, 2)