From e3f387b7fbaad4a69dad1f0e80785f88b40344b0 Mon Sep 17 00:00:00 2001
From: Nicholas Molnar <65710+neekolas@users.noreply.github.com>
Date: Tue, 20 Aug 2024 22:03:49 -0700
Subject: [PATCH] Fix broken tests

---
 xmtp_mls/src/groups/mod.rs       | 13 +++++++------
 xmtp_mls/src/identity_updates.rs | 32 ++++----------------------------
 2 files changed, 11 insertions(+), 34 deletions(-)

diff --git a/xmtp_mls/src/groups/mod.rs b/xmtp_mls/src/groups/mod.rs
index fd2bf327a..b842aa2ca 100644
--- a/xmtp_mls/src/groups/mod.rs
+++ b/xmtp_mls/src/groups/mod.rs
@@ -2918,16 +2918,17 @@ mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn ensure_removed_after_revoke() {
-        let wallet = generate_local_wallet();
-        let alix1 = ClientBuilder::new_test_client(&wallet).await;
-        let alix2 = ClientBuilder::new_test_client(&wallet).await;
-        let bo = ClientBuilder::new_test_client(&wallet).await;
+        let alix_wallet = generate_local_wallet();
+        let bo_wallet = generate_local_wallet();
+        let alix1 = ClientBuilder::new_test_client(&alix_wallet).await;
+        let alix2 = ClientBuilder::new_test_client(&alix_wallet).await;
+        let bo = ClientBuilder::new_test_client(&bo_wallet).await;
 
         let alix_group = alix1
             .create_group(None, GroupMetadataOptions::default())
             .unwrap();
         alix_group
-            .add_members(&alix1, vec![wallet.get_address()])
+            .add_members(&alix1, vec![bo_wallet.get_address()])
             .await
             .unwrap();
         let bo_group = receive_group_invite(&bo).await;
@@ -2943,7 +2944,7 @@ mod tests {
             .await
             .unwrap();
 
-        sign_with_wallet(&wallet, &mut revoke_installation_request).await;
+        sign_with_wallet(&alix_wallet, &mut revoke_installation_request).await;
         alix1
             .apply_signature_request(revoke_installation_request)
             .await
diff --git a/xmtp_mls/src/identity_updates.rs b/xmtp_mls/src/identity_updates.rs
index 87403fac8..817f40143 100644
--- a/xmtp_mls/src/identity_updates.rs
+++ b/xmtp_mls/src/identity_updates.rs
@@ -225,6 +225,7 @@ where
         existing_wallet_address: String,
         new_wallet_address: String,
     ) -> Result<SignatureRequest, ClientError> {
+        log::info!("Associating new wallet with inbox_id {}", self.inbox_id());
         let inbox_id = self.inbox_id();
         let builder = SignatureRequestBuilder::new(inbox_id);
 
@@ -518,19 +519,6 @@ pub(crate) mod tests {
         let wallet_2_address = wallet_2.get_address();
         let client = ClientBuilder::new_test_client(&wallet).await;
 
-        let mut signature_request: SignatureRequest = client
-            .create_inbox(wallet_address.clone(), None)
-            .await
-            .unwrap();
-        let inbox_id = signature_request.inbox_id();
-
-        sign_with_wallet(&wallet, &mut signature_request).await;
-
-        client
-            .apply_signature_request(signature_request)
-            .await
-            .unwrap();
-
         let mut add_association_request = client
             .associate_wallet(wallet_address.clone(), wallet_2_address.clone())
             .unwrap();
@@ -543,7 +531,7 @@ pub(crate) mod tests {
             .await
             .unwrap();
 
-        let association_state = get_association_state(&client, inbox_id.clone()).await;
+        let association_state = get_association_state(&client, client.inbox_id()).await;
 
         assert_eq!(association_state.members().len(), 3);
         assert_eq!(association_state.recovery_address(), &wallet_address);
@@ -558,19 +546,7 @@ pub(crate) mod tests {
         let wallet_address = wallet.get_address();
         let wallet_2_address = wallet_2.get_address();
         let client = ClientBuilder::new_test_client(&wallet).await;
-
-        let mut signature_request: SignatureRequest = client
-            .create_inbox(wallet_address.clone(), None)
-            .await
-            .unwrap();
-        let inbox_id = signature_request.inbox_id();
-
-        sign_with_wallet(&wallet, &mut signature_request).await;
-
-        client
-            .apply_signature_request(signature_request)
-            .await
-            .unwrap();
+        let inbox_id = client.inbox_id();
 
         get_association_state(&client, inbox_id.clone()).await;
 
@@ -772,7 +748,7 @@ pub(crate) mod tests {
         // Now revoke the second wallet
 
         let mut revoke_signature_request = client
-            .revoke_wallets(vec![second_wallet.get_address().into()])
+            .revoke_wallets(vec![second_wallet.get_address()])
             .await
             .unwrap();
         sign_with_wallet(&recovery_wallet, &mut revoke_signature_request).await;