From 1b8033426de7e88020811407456bb4dbb138679f Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Fri, 3 Nov 2023 20:41:22 -0400
Subject: [PATCH] Update processor docker tests

---
 tests/processor/src/tests/batch.rs   | 13 +++++++++----
 tests/processor/src/tests/key_gen.rs | 15 +++++++++------
 tests/processor/src/tests/send.rs    | 10 ++++++----
 3 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/tests/processor/src/tests/batch.rs b/tests/processor/src/tests/batch.rs
index ba2fbabf3..d41b767a7 100644
--- 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
index 903b2efac..f97eb2271 100644
--- 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
index 158262791..869180627 100644
--- 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"),
       }