Update processor docker tests

serai-dex · Nov 4, 2023 · 1b80334 · 1b80334
1 parent 63c6f83
commit 1b80334
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 14 deletions.
diff --git a/tests/processor/src/tests/batch.rs b/tests/processor/src/tests/batch.rs
@@ -48,7 +48,7 @@ pub(crate) async fn recv_batch_preprocesses(
         messages::coordinator::ProcessorMessage::BatchPreprocess {
           id: this_id,
           block: this_block,
-          preprocess,
+          preprocesses: mut these_preprocesses,
         },
       ) => {
         if id.is_none() {
@@ -60,7 +60,8 @@ pub(crate) async fn recv_batch_preprocesses(
         assert_eq!(&this_id, id.as_ref().unwrap());
         assert_eq!(&this_block, block.as_ref().unwrap());
 
-        preprocesses.insert(i, preprocess);
+        assert_eq!(these_preprocesses.len(), 1);
+        preprocesses.insert(i, these_preprocesses.swap_remove(0));
       }
       _ => panic!("processor didn't send batch preprocess"),
     }
@@ -107,10 +108,14 @@ pub(crate) async fn sign_batch(
     if preprocesses.contains_key(&i) {
       match coordinator.recv_message().await {
         messages::ProcessorMessage::Coordinator(
-          messages::coordinator::ProcessorMessage::BatchShare { id: this_id, share },
+          messages::coordinator::ProcessorMessage::BatchShare {
+            id: this_id,
+            shares: mut these_shares,
+          },
         ) => {
           assert_eq!(&this_id, &id);
-          shares.insert(i, share);
+          assert_eq!(these_shares.len(), 1);
+          shares.insert(i, these_shares.swap_remove(0));
         }
         _ => panic!("processor didn't send batch share"),
       }

diff --git a/tests/processor/src/tests/key_gen.rs b/tests/processor/src/tests/key_gen.rs
@@ -46,14 +46,16 @@ pub(crate) async fn key_gen(coordinators: &mut [Coordinator], network: NetworkId
         participant,
       )
       .unwrap(),
+      shares: 1,
     },
     |participant, msg| match msg {
       messages::key_gen::ProcessorMessage::Commitments {
         id: this_id,
-        commitments: these_commitments,
+        commitments: mut these_commitments,
       } => {
         assert_eq!(this_id, id);
-        commitments.insert(participant, these_commitments);
+        assert_eq!(these_commitments.len(), 1);
+        commitments.insert(participant, these_commitments.swap_remove(0));
       }
       _ => panic!("processor didn't return Commitments in response to GenerateKey"),
     },
@@ -69,9 +71,10 @@ pub(crate) async fn key_gen(coordinators: &mut [Coordinator], network: NetworkId
       commitments: clone_without(&commitments, &participant),
     },
     |participant, msg| match msg {
-      messages::key_gen::ProcessorMessage::Shares { id: this_id, shares: these_shares } => {
+      messages::key_gen::ProcessorMessage::Shares { id: this_id, shares: mut these_shares } => {
         assert_eq!(this_id, id);
-        shares.insert(participant, these_shares);
+        assert_eq!(these_shares.len(), 1);
+        shares.insert(participant, these_shares.swap_remove(0));
       }
       _ => panic!("processor didn't return Shares in response to GenerateKey"),
     },
@@ -85,12 +88,12 @@ pub(crate) async fn key_gen(coordinators: &mut [Coordinator], network: NetworkId
     coordinators,
     |participant| messages::key_gen::CoordinatorMessage::Shares {
       id,
-      shares: shares
+      shares: vec![shares
         .iter()
         .filter_map(|(this_participant, shares)| {
           shares.get(&participant).cloned().map(|share| (*this_participant, share))
         })
-        .collect(),
+        .collect()],
     },
     |_, msg| match msg {
       messages::key_gen::ProcessorMessage::GeneratedKeyPair {

diff --git a/tests/processor/src/tests/send.rs b/tests/processor/src/tests/send.rs
@@ -30,7 +30,7 @@ pub(crate) async fn recv_sign_preprocesses(
     match msg {
       messages::ProcessorMessage::Sign(messages::sign::ProcessorMessage::Preprocess {
         id: this_id,
-        preprocess,
+        preprocesses: mut these_preprocesses,
       }) => {
         if id.is_none() {
           assert_eq!(&this_id.key, &key);
@@ -39,7 +39,8 @@ pub(crate) async fn recv_sign_preprocesses(
         }
         assert_eq!(&this_id, id.as_ref().unwrap());
 
-        preprocesses.insert(i, preprocess);
+        assert_eq!(these_preprocesses.len(), 1);
+        preprocesses.insert(i, these_preprocesses.swap_remove(0));
       }
       _ => panic!("processor didn't send sign preprocess"),
     }
@@ -87,10 +88,11 @@ pub(crate) async fn sign_tx(
       match coordinator.recv_message().await {
         messages::ProcessorMessage::Sign(messages::sign::ProcessorMessage::Share {
           id: this_id,
-          share,
+          shares: mut these_shares,
         }) => {
           assert_eq!(&this_id, &id);
-          shares.insert(i, share);
+          assert_eq!(these_shares.len(), 1);
+          shares.insert(i, these_shares.swap_remove(0));
         }
         _ => panic!("processor didn't send TX shares"),
       }