From 3bc28d7756567b4d68b262bf51cdeb53f61836fc Mon Sep 17 00:00:00 2001
From: Piotr Roslaniec <p.roslaniec@gmail.com>
Date: Wed, 5 Jul 2023 13:40:51 +0200
Subject: [PATCH] update serialization tests where possible

---
 ferveo-python/ferveo/__init__.pyi        |  3 +++
 ferveo-python/test/test_serialization.py | 21 +++++++++++++++------
 ferveo/src/api.rs                        |  8 ++++++++
 ferveo/src/bindings_python.rs            | 13 +------------
 4 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/ferveo-python/ferveo/__init__.pyi b/ferveo-python/ferveo/__init__.pyi
index 170e98b0..b648301e 100644
--- a/ferveo-python/ferveo/__init__.pyi
+++ b/ferveo-python/ferveo/__init__.pyi
@@ -36,6 +36,9 @@ class FerveoPublicKey:
     def __hash__(self) -> int:
         ...
 
+    def __richcmp__(self, other: FerveoPublicKey, op: int) -> bool:
+        ...
+
 
 class Validator:
 
diff --git a/ferveo-python/test/test_serialization.py b/ferveo-python/test/test_serialization.py
index 6633ab68..30ba0ee9 100644
--- a/ferveo-python/test/test_serialization.py
+++ b/ferveo-python/test/test_serialization.py
@@ -4,6 +4,7 @@
     Dkg,
     DkgPublicKey,
     FerveoPublicKey,
+    SharedSecret,
 )
 
 
@@ -35,7 +36,8 @@ def make_dkg_public_key():
 
 
 def make_shared_secret():
-    # TODO: implement this
+    # TODO: Implement this
+    # SharedSecret.from_bytes(os.urandom(584))
     pass
 
 
@@ -44,27 +46,34 @@ def make_pk():
 
 
 # def test_shared_secret_serialization():
-#     shared_secret = create_shared_secret_instance()
+#     shared_secret = make_shared_secret()
 #     serialized = bytes(shared_secret)
 #     deserialized = SharedSecret.from_bytes(serialized)
-#     TODO: Implement comparison
-#     assert shared_secret == deserialized
+#     # TODO: Implement __richcmp__
+#     # assert shared_secret == deserialized
+#     assert serialized == bytes(deserialized)
 
 def test_keypair_serialization():
     keypair = Keypair.random()
     serialized = bytes(keypair)
     deserialized = Keypair.from_bytes(serialized)
-    # TODO: Implement comparison
-    # assert keypair == deserialized
+    # TODO: Implement __richcmp__
+    # assert serialized == deserialized
+    assert serialized == bytes(deserialized)
 
 
 def test_dkg_public_key_serialization():
     dkg_pk = make_dkg_public_key()
     serialized = bytes(dkg_pk)
+    deserialized = DkgPublicKey.from_bytes(serialized)
+    # TODO: Implement __richcmp__
+    assert serialized == bytes(deserialized)
     assert len(serialized) == DkgPublicKey.serialized_size()
 
 
 def test_public_key_serialization():
     pk = make_pk()
     serialized = bytes(pk)
+    deserialized = FerveoPublicKey.from_bytes(serialized)
+    assert pk == deserialized
     assert len(serialized) == FerveoPublicKey.serialized_size()
diff --git a/ferveo/src/api.rs b/ferveo/src/api.rs
index 99c5af02..398d4582 100644
--- a/ferveo/src/api.rs
+++ b/ferveo/src/api.rs
@@ -362,6 +362,14 @@ mod test_ferveo_api {
         (messages, validators, validator_keypairs)
     }
 
+    #[test]
+    fn test_dkg_pk_serialization() {
+        let dkg_pk = DkgPublicKey::random();
+        let serialized = dkg_pk.to_bytes().unwrap();
+        let deserialized = DkgPublicKey::from_bytes(&serialized).unwrap();
+        assert_eq!(dkg_pk, deserialized);
+    }
+
     #[test]
     fn test_server_api_tdec_precomputed() {
         let rng = &mut StdRng::seed_from_u64(0);
diff --git a/ferveo/src/bindings_python.rs b/ferveo/src/bindings_python.rs
index 05756164..614b3608 100644
--- a/ferveo/src/bindings_python.rs
+++ b/ferveo/src/bindings_python.rs
@@ -172,17 +172,6 @@ where
     }
 }
 
-// TODO: Consider implementing macros to generate following methods
-
-// fn __richcmp__(&self, other: &Self, op: CompareOp) -> PyResult<bool> {
-//     richcmp(self, other, op)
-// }
-
-// fn __hash__(&self) -> PyResult<isize> {
-//     let bytes = self.0.to_bytes()?;
-//     hash(stringify!($struct_name), &bytes)
-// }
-
 macro_rules! generate_bytes_serialization {
     ($struct_name:ident) => {
         #[pymethods]
@@ -326,6 +315,7 @@ generate_boxed_bytes_serialization!(FerveoPublicKey, InnerPublicKey);
 
 #[pymethods]
 impl FerveoPublicKey {
+    // We implement `__richcmp__` because FerveoPublicKeys must be sortable in some cases
     fn __richcmp__(&self, other: &Self, op: CompareOp) -> PyResult<bool> {
         richcmp(self, other, op)
     }
@@ -857,7 +847,6 @@ mod test_ferveo_python {
 
         let shared_secret = combine_decryption_shares_simple(decryption_shares);
 
-        // TODO: Fails because of a bad shared secret
         let plaintext =
             decrypt_with_shared_secret(&ciphertext, aad, &shared_secret)
                 .unwrap();