From da7a3c41873ede0ba83546e97c7a3181c05ba42a Mon Sep 17 00:00:00 2001
From: Dakota Brink <git@kota.is>
Date: Tue, 1 Oct 2024 11:12:32 -0400
Subject: [PATCH] restore tests too

---
 mls_validation_service/src/handlers.rs | 101 +++++++++++++++++++++++--
 1 file changed, 96 insertions(+), 5 deletions(-)

diff --git a/mls_validation_service/src/handlers.rs b/mls_validation_service/src/handlers.rs
index dc3cb6c28..37f4f0976 100644
--- a/mls_validation_service/src/handlers.rs
+++ b/mls_validation_service/src/handlers.rs
@@ -463,18 +463,24 @@ mod tests {
     };
     use openmls_basic_credential::SignatureKeyPair;
     use openmls_rust_crypto::OpenMlsRustCrypto;
-    use xmtp_id::associations::{
-        generate_inbox_id,
-        test_utils::{rand_string, rand_u64, MockSmartContractSignatureVerifier},
-        unverified::{UnverifiedAction, UnverifiedIdentityUpdate},
+    use xmtp_id::{
+        associations::{
+            generate_inbox_id,
+            test_utils::{rand_string, rand_u64, MockSmartContractSignatureVerifier},
+            unverified::{UnverifiedAction, UnverifiedIdentityUpdate},
+        },
+        InboxOwner,
     };
-    use xmtp_mls::configuration::CIPHERSUITE;
+    use xmtp_mls::{configuration::CIPHERSUITE, credential::Credential};
     use xmtp_proto::xmtp::{
         identity::associations::IdentityUpdate as IdentityUpdateProto,
         identity::MlsCredential as InboxIdMlsCredential,
+        mls::message_contents::MlsCredential as CredentialProto,
         mls_validation::v1::validate_key_packages_request::KeyPackage as KeyPackageProtoWrapper,
     };
 
+    use prost::Message;
+
     use super::*;
 
     impl Default for ValidationService {
@@ -483,6 +489,25 @@ mod tests {
         }
     }
 
+    fn generate_identity() -> (Vec<u8>, SignatureKeyPair, String) {
+        let rng = &mut rand::thread_rng();
+        let wallet = LocalWallet::new(rng);
+        let signature_key_pair = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm()).unwrap();
+
+        let _pub_key = signature_key_pair.public();
+        let account_address = wallet.get_address();
+
+        let credential =
+            Credential::create(&signature_key_pair, &wallet).expect("failed to create credential");
+        let credential_proto: CredentialProto = credential.into();
+
+        (
+            credential_proto.encode_to_vec(),
+            signature_key_pair,
+            account_address,
+        )
+    }
+
     async fn generate_inbox_id_credential() -> (String, SigningKey) {
         let signing_key = SigningKey::generate(&mut rand::thread_rng());
 
@@ -533,6 +558,72 @@ mod tests {
         )
     }
 
+    #[tokio::test]
+    async fn test_validate_key_packages_happy_path() {
+        let (identity, keypair, account_address) = generate_identity();
+
+        let credential: OpenMlsCredential = BasicCredential::new(identity).into();
+        let credential_with_key = CredentialWithKey {
+            credential,
+            signature_key: keypair.to_public_vec().into(),
+        };
+
+        let key_package_bytes = build_key_package_bytes(
+            &keypair,
+            &credential_with_key,
+            Some(account_address.clone()),
+        );
+        let request = ValidateKeyPackagesRequest {
+            key_packages: vec![KeyPackageProtoWrapper {
+                key_package_bytes_tls_serialized: key_package_bytes,
+                is_inbox_id_credential: false,
+            }],
+        };
+
+        let res = ValidationService::default()
+            .validate_key_packages(Request::new(request))
+            .await
+            .unwrap();
+
+        let first_response = &res.into_inner().responses[0];
+        assert_eq!(first_response.installation_id, keypair.public());
+        assert_eq!(first_response.account_address, account_address);
+        assert!(!first_response.credential_identity_bytes.is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_validate_key_packages_fail() {
+        let (identity, keypair, account_address) = generate_identity();
+        let (_, other_keypair, _) = generate_identity();
+
+        let credential: OpenMlsCredential = BasicCredential::new(identity).into();
+        let credential_with_key = CredentialWithKey {
+            credential,
+            // Use the wrong signature key to make the validation fail
+            signature_key: other_keypair.to_public_vec().into(),
+        };
+
+        let key_package_bytes =
+            build_key_package_bytes(&keypair, &credential_with_key, Some(account_address));
+
+        let request = ValidateKeyPackagesRequest {
+            key_packages: vec![KeyPackageProtoWrapper {
+                key_package_bytes_tls_serialized: key_package_bytes,
+                is_inbox_id_credential: false,
+            }],
+        };
+
+        let res = ValidationService::default()
+            .validate_key_packages(Request::new(request))
+            .await
+            .unwrap();
+
+        let first_response = &res.into_inner().responses[0];
+
+        assert!(!first_response.is_ok);
+        assert_eq!(first_response.account_address, "".to_string());
+    }
+
     // this test will panic until signature recovery is added
     // and `MockSignature` is updated with signatures that can be recovered
     #[tokio::test]