Repository: tangle-network/gadget (language: rust)

Issue:
DMSM and DFFT Integration tests (using test-gadget)



Discussion / hints:

Fix patch:
diff --git a/Cargo.toml b/Cargo.toml
index 5a9b42d44..e117c967b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,9 @@ sp-runtime = { git = "https://github.com/paritytech/substrate", branch = "polkad
 sp-api = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v1.0.0", default-features = false }
 
 mpc-net = { git = "https://github.com/webb-tools/zk-SaaS/" }
+dist-primitives = { git = "https://github.com/webb-tools/zk-SaaS/" }
+secret-sharing = { git = "https://github.com/webb-tools/zk-SaaS/" }
+
 tokio-rustls = "0.24.1"
 tokio = "1.32.0"
 bincode2 = "2"
diff --git a/gadget-core/src/gadget/manager.rs b/gadget-core/src/gadget/manager.rs
index 3c8e5e748..ae682a5de 100644
--- a/gadget-core/src/gadget/manager.rs
+++ b/gadget-core/src/gadget/manager.rs
@@ -1,10 +1,11 @@
 use async_trait::async_trait;
 use std::error::Error;
+use std::fmt::{Display, Formatter};
 use std::future::Future;
 use std::pin::Pin;
 
 pub struct GadgetManager<'a> {
-    gadget: Pin<Box<dyn Future<Output = Result<(), GadgetError>> + 'a>>,
+    gadget: Pin<Box<dyn Future<Output = Result<(), GadgetError>> + Send + 'a>>,
 }
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -14,12 +15,20 @@ pub enum GadgetError {
     ProtocolMessageStreamEnded,
 }
 
+impl Display for GadgetError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Debug::fmt(self, f)
+    }
+}
+
+impl std::error::Error for GadgetError {}
+
 #[async_trait]
-pub trait AbstractGadget: Send {
+pub trait AbstractGadget: Send + Sync {
     type FinalityNotification: Send;
     type BlockImportNotification: Send;
     type ProtocolMessage: Send;
-    type Error: Error;
+    type Error: Error + Send;
 
     async fn get_next_finality_notification(&self) -> Option<Self::FinalityNotification>;
     async fn get_next_block_import_notification(&self) -> Option<Self::BlockImportNotification>;
diff --git a/gadget-core/src/job_manager.rs b/gadget-core/src/job_manager.rs
index 3e6bc7c4a..10315492a 100644
--- a/gadget-core/src/job_manager.rs
+++ b/gadget-core/src/job_manager.rs
@@ -426,7 +426,7 @@ impl<WM: WorkManagerInterface> ProtocolWorkManager<WM> {
         for task in lock.enqueued_tasks.iter() {
             if should_deliver(task, &msg, message_task_hash) {
                 self.utility.debug(format!(
-                    "Message is for this ENQUEUED signing execution in session: {}",
+                    "Message is for this ENQUEUED execution in session: {}",
                     task.handle.session_id()
                 ));
                 if let Err(err) = task.handle.deliver_message(msg) {
@@ -443,7 +443,7 @@ impl<WM: WorkManagerInterface> ProtocolWorkManager<WM> {
         for task in lock.active_tasks.iter() {
             if should_deliver(task, &msg, message_task_hash) {
                 self.utility.debug(format!(
-                    "Message is for this signing CURRENT execution in session: {}",
+                    "Message is for this CURRENT execution in session: {}",
                     task.handle.session_id()
                 ));
                 if let Err(err) = task.handle.deliver_message(msg) {
@@ -533,7 +533,7 @@ impl<WM: WorkManagerInterface> Hash for Job<WM> {
 impl<WM: WorkManagerInterface> Drop for Job<WM> {
     fn drop(&mut self) {
         self.utility.debug(format!(
-            "Will remove job {:?} from currently_signing_proposals",
+            "Will remove job {:?} from JobManager",
             hex::encode(self.task_hash)
         ));
         let _ = self.handle.shutdown(ShutdownReason::DropCode);
diff --git a/webb-gadget/src/gadget/mod.rs b/webb-gadget/src/gadget/mod.rs
index 0947bf0cf..682330bdd 100644
--- a/webb-gadget/src/gadget/mod.rs
+++ b/webb-gadget/src/gadget/mod.rs
@@ -18,7 +18,7 @@ pub mod network;
 pub mod work_manager;
 
 /// Used as a module to place inside the SubstrateGadget
-pub struct WebbGadget<B, C, BE, N, M> {
+pub struct WebbModule<B, C, BE, N, M> {
     #[allow(dead_code)]
     network: N,
     module: M,
@@ -32,7 +32,7 @@ const MAX_ACTIVE_TASKS: usize = 4;
 const MAX_PENDING_TASKS: usize = 4;
 
 impl<C: Client<B, BE>, B: Block, BE: Backend<B>, N: Network, M: WebbGadgetModule<B>>
-    WebbGadget<B, C, BE, N, M>
+    WebbModule<B, C, BE, N, M>
 {
     pub fn new(mut network: N, module: M, now: Option<u64>) -> Self {
         let clock = Arc::new(RwLock::new(now));
@@ -48,7 +48,7 @@ impl<C: Client<B, BE>, B: Block, BE: Backend<B>, N: Network, M: WebbGadgetModule
             PollMethod::Interval { millis: 200 },
         );
 
-        WebbGadget {
+        WebbModule {
             module,
             network,
             job_manager,
@@ -61,7 +61,7 @@ impl<C: Client<B, BE>, B: Block, BE: Backend<B>, N: Network, M: WebbGadgetModule
 
 #[async_trait]
 impl<C: Client<B, BE>, B: Block, BE: Backend<B>, N: Network, M: WebbGadgetModule<B>>
-    SubstrateGadgetModule for WebbGadget<B, C, BE, N, M>
+    SubstrateGadgetModule for WebbModule<B, C, BE, N, M>
 {
     type Error = Error;
     type ProtocolMessage = GadgetProtocolMessage;
diff --git a/webb-gadget/src/lib.rs b/webb-gadget/src/lib.rs
index 0105e0229..f31a7d98f 100644
--- a/webb-gadget/src/lib.rs
+++ b/webb-gadget/src/lib.rs
@@ -1,5 +1,5 @@
 use crate::gadget::network::Network;
-use crate::gadget::{WebbGadget, WebbGadgetModule};
+use crate::gadget::{WebbGadgetModule, WebbModule};
 use gadget_core::gadget::manager::{GadgetError, GadgetManager};
 use gadget_core::gadget::substrate::{Client, SubstrateGadget};
 use gadget_core::job_manager::WorkManagerError;
@@ -35,9 +35,9 @@ pub async fn run<C: Client<B, BE>, B: Block, BE: Backend<B>, N: Network, M: Webb
     client: C,
 ) -> Result<(), Error> {
     let now = None;
-    let webb_gadget = WebbGadget::new(network, module, now);
+    let webb_module = WebbModule::new(network, module, now);
     // Plug the module into the substrate gadget to interface the WebbGadget with Substrate
-    let substrate_gadget = SubstrateGadget::new(client, webb_gadget);
+    let substrate_gadget = SubstrateGadget::new(client, webb_module);
 
     // Run the GadgetManager to execute the substrate gadget
     GadgetManager::new(substrate_gadget)
diff --git a/zk-gadget/src/lib.rs b/zk-gadget/src/lib.rs
index c128b9bdb..b95c450b8 100644
--- a/zk-gadget/src/lib.rs
+++ b/zk-gadget/src/lib.rs
@@ -1,5 +1,5 @@
 use crate::module::ZkModule;
-use crate::network::RegistantId;
+use crate::network::{RegistantId, ZkNetworkService};
 use gadget_core::gadget::substrate::Client;
 use gadget_core::{Backend, Block};
 use mpc_net::prod::RustlsCertificate;
@@ -24,13 +24,21 @@ pub async fn run<C: Client<B, BE>, B: Block, BE: Backend<B>>(
     client: C,
 ) -> Result<(), Error> {
     // Create the zk gadget module
+    let network = create_zk_network(&config).await?;
+
+    let zk_module = ZkModule {}; // TODO: proper implementation
+
+    // Plug the module into the webb gadget
+    webb_gadget::run(network, zk_module, client).await
+}
+
+async fn create_zk_network(config: &ZkGadgetConfig) -> Result<ZkNetworkService, Error> {
     let our_identity = RustlsCertificate {
-        cert: Certificate(config.public_identity_der),
-        private_key: PrivateKey(config.private_identity_der),
+        cert: Certificate(config.public_identity_der.clone()),
+        private_key: PrivateKey(config.private_identity_der.clone()),
     };
-
-    let network = if let Some(addr) = &config.king_bind_addr {
-        network::ZkNetworkService::new_king(*addr, our_identity).await?
+    if let Some(addr) = &config.king_bind_addr {
+        ZkNetworkService::new_king(*addr, our_identity).await
     } else {
         let king_addr = config
             .client_only_king_addr
@@ -41,18 +49,13 @@ pub async fn run<C: Client<B, BE>, B: Block, BE: Backend<B>>(
             .add(&Certificate(
                 config
                     .client_only_king_public_identity_der
+                    .clone()
                     .expect("The client must specify the identity of the king"),
             ))
             .map_err(|err| Error::InitError {
                 err: err.to_string(),
             })?;
 
-        network::ZkNetworkService::new_client(king_addr, config.id, our_identity, king_certs)
-            .await?
-    };
-
-    let zk_module = ZkModule {}; // TODO: proper implementation
-
-    // Plug the module into the webb gadget
-    webb_gadget::run(network, zk_module, client).await
+        ZkNetworkService::new_client(king_addr, config.id, our_identity, king_certs).await
+    }
 }


Test patch:
diff --git a/test-gadget/Cargo.toml b/test-gadget/Cargo.toml
index e44de903c..a200d246e 100644
--- a/test-gadget/Cargo.toml
+++ b/test-gadget/Cargo.toml
@@ -17,4 +17,18 @@ async-trait = { workspace = true }
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] }
-tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
\ No newline at end of file
+tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
+
+# zk deps
+mpc-net = { workspace = true }
+dist-primitives = { workspace = true }
+secret-sharing = { workspace = true }
+bytes = { version = "1.5.0", features = ["serde"]}
+ark-std = {version = "0.4.0", default-features = false, features = [ "print-trace", "std" ]}
+ark-ec = {version = "0.4.0", default-features = false}
+ark-bls12-377 = { version = "0.4.0", default-features = false, features = ["curve"] }
+ark-ff = {version = "0.4.0", default-features = false}
+ark-poly = {version = "0.4.0", default-features = false}
+bincode2 = "2.0.1"
+hex = "0.4.3"
+md5 = "0.7.0"
diff --git a/test-gadget/src/gadget.rs b/test-gadget/src/gadget.rs
index cde943c9d..4653ecabc 100644
--- a/test-gadget/src/gadget.rs
+++ b/test-gadget/src/gadget.rs
@@ -13,21 +13,23 @@ use std::sync::Arc;
 use tokio::sync::Mutex;
 
 /// An AbstractGadget endowed with a WorkerManager, a fake blockchain that delivers FinalityNotifications to the gadgets, and a TestProtocolMessage stream
-pub struct TestGadget {
+pub struct TestGadget<B> {
     job_manager: ProtocolWorkManager<TestWorkManager>,
     blockchain_connection: Mutex<tokio::sync::broadcast::Receiver<TestFinalityNotification>>,
     network_connection: Mutex<tokio::sync::mpsc::UnboundedReceiver<TestProtocolMessage>>,
     // Specifies at which blocks we should start a job on
     run_test_at: Arc<Vec<u64>>,
     clock: Arc<RwLock<u64>>,
-    async_protocol_generator: Box<dyn AsyncProtocolGenerator>,
+    test_bundle: B,
+    async_protocol_generator: Box<dyn AsyncProtocolGenerator<B>>,
 }
 
-impl TestGadget {
-    pub fn new<T: AsyncProtocolGenerator + 'static>(
+impl<B: Send + Sync + Clone + 'static> TestGadget<B> {
+    pub fn new<T: AsyncProtocolGenerator<B> + 'static>(
         blockchain_connection: tokio::sync::broadcast::Receiver<TestFinalityNotification>,
         network_connection: tokio::sync::mpsc::UnboundedReceiver<TestProtocolMessage>,
         run_test_at: Arc<Vec<u64>>,
+        test_bundle: B,
         async_protocol_generator: T,
     ) -> Self {
         let clock = Arc::new(RwLock::new(0));
@@ -43,6 +45,7 @@ impl TestGadget {
             blockchain_connection: Mutex::new(blockchain_connection),
             network_connection: Mutex::new(network_connection),
             run_test_at,
+            test_bundle,
             async_protocol_generator: Box::new(async_protocol_generator),
             clock,
         }
@@ -56,7 +59,7 @@ pub struct TestFinalityNotification {
 }
 
 #[async_trait]
-impl AbstractGadget for TestGadget {
+impl<B: Send + Sync + Clone + 'static> AbstractGadget for TestGadget<B> {
     type FinalityNotification = TestFinalityNotification;
     type BlockImportNotification = ();
     type ProtocolMessage = TestProtocolMessage;
@@ -83,6 +86,7 @@ impl AbstractGadget for TestGadget {
         let session_id = notification.session_id;
         *self.clock.write() = now;
 
+        // Determine if we need to run a task
         if self.run_test_at.contains(&now) {
             log::info!("Running test at block {now}");
             let task_hash = now.to_be_bytes();
@@ -92,6 +96,7 @@ impl AbstractGadget for TestGadget {
                 now,
                 ssid,
                 task_hash,
+                self.test_bundle.clone(),
                 &*self.async_protocol_generator,
             );
             self.job_manager
@@ -128,12 +133,13 @@ impl AbstractGadget for TestGadget {
     }
 }
 
-fn create_test_async_protocol(
+fn create_test_async_protocol<B: Send + Sync + 'static>(
     session_id: <TestWorkManager as WorkManagerInterface>::SessionID,
     now: <TestWorkManager as WorkManagerInterface>::Clock,
     ssid: <TestWorkManager as WorkManagerInterface>::SSID,
     task_id: <TestWorkManager as WorkManagerInterface>::TaskID,
-    proto_gen: &dyn AsyncProtocolGenerator,
+    test_bundle: B,
+    proto_gen: &dyn AsyncProtocolGenerator<B>,
 ) -> (TestProtocolRemote, Pin<Box<dyn SendFuture<'static, ()>>>) {
     let is_done = Arc::new(AtomicBool::new(false));
     let (to_async_protocol, protocol_message_rx) = tokio::sync::mpsc::unbounded_channel();
@@ -149,6 +155,7 @@ fn create_test_async_protocol(
         associated_ssid: ssid,
         associated_session_id: session_id,
         associated_task_id: task_id,
+        test_bundle,
     };
 
     let remote = TestProtocolRemote {
diff --git a/test-gadget/src/lib.rs b/test-gadget/src/lib.rs
index 8deb56a6a..995c7df0a 100644
--- a/test-gadget/src/lib.rs
+++ b/test-gadget/src/lib.rs
@@ -1,6 +1,110 @@
+use crate::error::TestError;
+use crate::gadget::TestGadget;
+use crate::message::UserID;
+use crate::test_network::InMemoryNetwork;
+use crate::work_manager::AsyncProtocolGenerator;
+use futures::stream::FuturesUnordered;
+use futures::TryStreamExt;
+use gadget_core::gadget::manager::GadgetManager;
+use std::error::Error;
+use std::sync::Arc;
+use std::time::Duration;
+
 pub mod blockchain;
 pub mod error;
 pub mod gadget;
 pub mod message;
 pub mod test_network;
 pub mod work_manager;
+
+pub async fn simulate_test<T: AsyncProtocolGenerator<TestBundle> + 'static + Clone>(
+    n_peers: usize,
+    n_blocks_per_session: u64,
+    block_duration: Duration,
+    run_tests_at: Vec<u64>,
+    async_proto_gen: T,
+) -> Result<(), Box<dyn Error>> {
+    let (network, recv_handles) = InMemoryNetwork::new(n_peers);
+    let (bc_tx, _bc_rx) = tokio::sync::broadcast::channel(5000);
+    let gadget_futures = FuturesUnordered::new();
+    let run_tests_at = Arc::new(run_tests_at);
+    let (count_finished_tx, mut count_finished_rx) = tokio::sync::mpsc::unbounded_channel();
+
+    for (party_id, network_handle) in recv_handles.into_iter().enumerate() {
+        // Create a TestGadget
+        let count_finished_tx = count_finished_tx.clone();
+        let network = network.clone();
+        let test_bundle = TestBundle {
+            count_finished_tx,
+            network,
+            n_peers,
+            party_id: party_id as UserID,
+        };
+
+        let gadget = TestGadget::new(
+            bc_tx.subscribe(),
+            network_handle,
+            run_tests_at.clone(),
+            test_bundle,
+            async_proto_gen.clone(),
+        );
+
+        gadget_futures.push(async move {
+            tokio::task::spawn(async move {
+                GadgetManager::new(gadget).await.map_err(|err| TestError {
+                    reason: format!("{err:?}"),
+                })
+            })
+            .await
+            .map_err(|err| TestError {
+                reason: format!("{err:?}"),
+            })?
+        });
+    }
+
+    let blockchain_future = blockchain::blockchain(block_duration, n_blocks_per_session, bc_tx);
+
+    let finished_future = async move {
+        let mut count_received = 0;
+        let expected_count = n_peers * run_tests_at.len();
+        while count_finished_rx.recv().await.is_some() {
+            count_received += 1;
+            log::info!("Received {} finished signals", count_received);
+            if count_received == expected_count {
+                return Ok::<(), TestError>(());
+            }
+        }
+
+        Err(TestError {
+            reason: "Didn't receive all finished signals".to_string(),
+        })
+    };
+
+    let gadgets_future = gadget_futures.try_collect::<Vec<_>>();
+
+    // The gadgets will run indefinitely if properly behaved, and the blockchain future will as well.
+    // The finished future will end when all gadgets have finished their async protocols properly
+    // Thus, select all three futures
+
+    tokio::select! {
+            res0 = blockchain_future => {
+                res0?;
+            },
+            res1 = gadgets_future => {
+                res1?;
+            },
+            res2 = finished_future => {
+                res2?;
+            }
+    }
+
+    Ok(())
+}
+
+#[derive(Clone)]
+pub struct TestBundle {
+    pub count_finished_tx: tokio::sync::mpsc::UnboundedSender<()>,
+    pub network: InMemoryNetwork,
+    pub n_peers: usize,
+    pub party_id: UserID,
+}
diff --git a/test-gadget/src/message.rs b/test-gadget/src/message.rs
index 5e0b96415..bca454d5a 100644
--- a/test-gadget/src/message.rs
+++ b/test-gadget/src/message.rs
@@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize};
 #[derive(Serialize, Deserialize, Clone)]
 pub struct TestProtocolMessage {
     pub payload: Vec<u8>,
-    pub from: u32,
-    pub to: Option<u32>,
+    pub from: UserID,
+    pub to: Option<UserID>,
     pub associated_block_id: <TestWorkManager as WorkManagerInterface>::Clock,
     pub associated_session_id: <TestWorkManager as WorkManagerInterface>::SessionID,
     pub associated_ssid: <TestWorkManager as WorkManagerInterface>::SSID,
diff --git a/test-gadget/src/work_manager.rs b/test-gadget/src/work_manager.rs
index 194b2e493..3ba8c882f 100644
--- a/test-gadget/src/work_manager.rs
+++ b/test-gadget/src/work_manager.rs
@@ -25,6 +25,7 @@ pub struct TestProtocolRemote {
 const ACCEPTABLE_BLOCK_TOLERANCE: u64 = 3;
 
 impl WorkManagerInterface for TestWorkManager {
+    // TODO: rename to retry_id or similar
     type SSID = u16;
     type Clock = u64;
     type ProtocolMessage = TestProtocolMessage;
@@ -115,7 +116,7 @@ impl ProtocolRemote<TestWorkManager> for TestProtocolRemote {
     }
 }
 
-pub struct TestAsyncProtocolParameters {
+pub struct TestAsyncProtocolParameters<B> {
     pub is_done: Arc<AtomicBool>,
     pub protocol_message_rx: tokio::sync::mpsc::UnboundedReceiver<
         <TestWorkManager as WorkManagerInterface>::ProtocolMessage,
@@ -126,13 +127,16 @@ pub struct TestAsyncProtocolParameters {
     pub associated_ssid: <TestWorkManager as WorkManagerInterface>::SSID,
     pub associated_session_id: <TestWorkManager as WorkManagerInterface>::SessionID,
     pub associated_task_id: <TestWorkManager as WorkManagerInterface>::TaskID,
+    pub test_bundle: B,
 }
 
-pub trait AsyncProtocolGenerator:
-    Send + Sync + Fn(TestAsyncProtocolParameters) -> Pin<Box<dyn SendFuture<'static, ()>>>
+pub trait AsyncProtocolGenerator<B>:
+    Send + Sync + Fn(TestAsyncProtocolParameters<B>) -> Pin<Box<dyn SendFuture<'static, ()>>>
 {
 }
-impl<T: Send + Sync + Fn(TestAsyncProtocolParameters) -> Pin<Box<dyn SendFuture<'static, ()>>>>
-    AsyncProtocolGenerator for T
+impl<
+        B: Send + Sync,
+        T: Send + Sync + Fn(TestAsyncProtocolParameters<B>) -> Pin<Box<dyn SendFuture<'static, ()>>>,
+    > AsyncProtocolGenerator<B> for T
 {
 }
diff --git a/test-gadget/tests/basic.rs b/test-gadget/tests/basic.rs
index 287a819fa..40bcbc2e8 100644
--- a/test-gadget/tests/basic.rs
+++ b/test-gadget/tests/basic.rs
@@ -1,18 +1,12 @@
 #[cfg(test)]
 mod tests {
-    use futures::stream::FuturesUnordered;
-    use futures::TryStreamExt;
-    use gadget_core::gadget::manager::GadgetManager;
     use gadget_core::job_manager::SendFuture;
     use std::error::Error;
     use std::pin::Pin;
-    use std::sync::Arc;
     use std::time::Duration;
-    use test_gadget::error::TestError;
-    use test_gadget::gadget::TestGadget;
-    use test_gadget::message::{TestProtocolMessage, UserID};
-    use test_gadget::test_network::InMemoryNetwork;
+    use test_gadget::message::TestProtocolMessage;
     use test_gadget::work_manager::TestAsyncProtocolParameters;
+    use test_gadget::TestBundle;
     use tracing_subscriber::fmt::SubscriberBuilder;
     use tracing_subscriber::util::SubscriberInitExt;
     use tracing_subscriber::EnvFilter;
@@ -29,88 +23,21 @@ mod tests {
     #[tokio::test]
     async fn test_basic_async_protocol() -> Result<(), Box<dyn Error>> {
         setup_log();
-        const N_PEERS: usize = 5;
-        const N_BLOCKS_PER_SESSION: u64 = 10;
-        const BLOCK_DURATION: Duration = Duration::from_millis(1000);
-
-        let (network, recv_handles) = InMemoryNetwork::new(N_PEERS);
-        let (bc_tx, _bc_rx) = tokio::sync::broadcast::channel(1000);
-        let run_tests_at = Arc::new(vec![0, 2, 4]);
-        let gadget_futures = FuturesUnordered::new();
-        let (count_finished_tx, mut count_finished_rx) = tokio::sync::mpsc::unbounded_channel();
-
-        for (party_id, network_handle) in recv_handles.into_iter().enumerate() {
-            // Create a TestGadget
-            let count_finished_tx = count_finished_tx.clone();
-            let network = network.clone();
-            let async_protocol_gen = move |params: TestAsyncProtocolParameters| {
-                async_proto_generator(
-                    params,
-                    count_finished_tx.clone(),
-                    network.clone(),
-                    N_PEERS,
-                    party_id as UserID,
-                )
-            };
-
-            let gadget = TestGadget::new(
-                bc_tx.subscribe(),
-                network_handle,
-                run_tests_at.clone(),
-                async_protocol_gen,
-            );
-
-            gadget_futures.push(GadgetManager::new(gadget));
-        }
-
-        let blockchain_future =
-            test_gadget::blockchain::blockchain(BLOCK_DURATION, N_BLOCKS_PER_SESSION, bc_tx);
-
-        let finished_future = async move {
-            let mut count_received = 0;
-            let expected_count = N_PEERS * run_tests_at.len();
-            while count_finished_rx.recv().await.is_some() {
-                count_received += 1;
-                log::info!("Received {} finished signals", count_received);
-                if count_received == expected_count {
-                    return Ok::<(), TestError>(());
-                }
-            }
-
-            Err(TestError {
-                reason: "Didn't receive all finished signals".to_string(),
-            })
-        };
-
-        let gadgets_future = gadget_futures.try_collect::<Vec<_>>();
-
-        // The gadgets will run indefinitely if properly behaved, and the blockchain future will as well.
-        // The finished future will end when all gadgets have finished their async protocols properly
-        // Thus, select all three futures
-
-        tokio::select! {
-            res0 = blockchain_future => {
-                res0?;
-            },
-            res1 = gadgets_future => {
-                res1.map_err(|err| TestError { reason: format!("{err:?}") })
-                    .map(|_| ())?;
-            },
-            res2 = finished_future => {
-                res2?;
-            }
-        }
-
-        Ok(())
+        test_gadget::simulate_test(
+            5,
+            10,
+            Duration::from_millis(1000),
+            vec![0, 2, 4],
+            async_proto_generator,
+        )
+        .await
     }
 
     fn async_proto_generator(
-        mut params: TestAsyncProtocolParameters,
-        on_finish: tokio::sync::mpsc::UnboundedSender<()>,
-        network: InMemoryNetwork,
-        n_peers: usize,
-        party_id: UserID,
+        mut params: TestAsyncProtocolParameters<TestBundle>,
     ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+        let n_peers = params.test_bundle.n_peers;
+        let party_id = params.test_bundle.party_id;
         Box::pin(async move {
             params
                 .start_rx
@@ -119,7 +46,9 @@ mod tests {
                 .await
                 .expect("Failed to start");
             // Broadcast a message to each peer
-            network
+            params
+                .test_bundle
+                .network
                 .broadcast(
                     party_id,
                     TestProtocolMessage {
@@ -145,7 +74,11 @@ mod tests {
             params
                 .is_done
                 .store(true, std::sync::atomic::Ordering::Relaxed);
-            on_finish.send(()).expect("Didn't send on_finish signal");
+            params
+                .test_bundle
+                .count_finished_tx
+                .send(())
+                .expect("Didn't send on_finish signal");
         })
     }
 }
diff --git a/test-gadget/tests/dfft.rs b/test-gadget/tests/dfft.rs
new file mode 100644
index 000000000..8afbc0df3
--- /dev/null
+++ b/test-gadget/tests/dfft.rs
@@ -0,0 +1,298 @@
+#[cfg(test)]
+mod tests {
+    use gadget_core::job_manager::SendFuture;
+    use std::collections::HashMap;
+    use std::error::Error;
+    use std::pin::Pin;
+    use test_gadget::work_manager::TestAsyncProtocolParameters;
+    use test_gadget::TestBundle;
+    use tracing_subscriber::fmt::SubscriberBuilder;
+    use tracing_subscriber::util::SubscriberInitExt;
+    use tracing_subscriber::EnvFilter;
+
+    use ark_bls12_377::Fr;
+    use ark_ff::{FftField, PrimeField};
+    use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
+    use async_trait::async_trait;
+    use bytes::Bytes;
+    use dist_primitives::{
+        channel::MpcSerNet,
+        dfft::{d_fft, fft_in_place_rearrange},
+        utils::pack::transpose,
+    };
+    use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
+    use secret_sharing::pss::PackedSharingParams;
+    use serde::{Deserialize, Serialize};
+    use test_gadget::message::TestProtocolMessage;
+    use test_gadget::test_network::InMemoryNetwork;
+    use tokio::sync::Mutex;
+
+    const L: usize = 2;
+    const N: usize = L * 4;
+    const M: usize = L * 32;
+
+    pub async fn d_fft_test<F: FftField + PrimeField, Net: MpcNet>(
+        pp: &PackedSharingParams<F>,
+        dom: &Radix2EvaluationDomain<F>,
+        net: &Net,
+    ) {
+        let mbyl: usize = dom.size() / pp.l;
+        // We apply FFT on this vector
+        // let mut x = vec![F::ONE; cd.m];
+        let mut x: Vec<F> = Vec::new();
+        for i in 0..dom.size() {
+            x.push(F::from(i as u64));
+        }
+
+        // Output to test against
+        let should_be_output = dom.fft(&x);
+
+        fft_in_place_rearrange(&mut x);
+        let mut pcoeff: Vec<Vec<F>> = Vec::new();
+        for i in 0..mbyl {
+            pcoeff.push(x.iter().skip(i).step_by(mbyl).cloned().collect::<Vec<_>>());
+            pp.pack_from_public_in_place(&mut pcoeff[i]);
+        }
+
+        let pcoeff_share = pcoeff
+            .iter()
+            .map(|x| x[net.party_id() as usize])
+            .collect::<Vec<_>>();
+
+        // Rearranging x
+
+        let peval_share = d_fft(
+            pcoeff_share,
+            false,
+            1,
+            false,
+            dom,
+            pp,
+            net,
+            MultiplexedStreamID::One,
+        )
+        .await
+        .unwrap();
+
+        // Send to king who reconstructs and checks the answer
+        let result = net
+            .send_to_king(&peval_share, MultiplexedStreamID::One)
+            .await
+            .unwrap();
+
+        if let Some(peval_shares) = result {
+            let peval_shares = transpose(peval_shares);
+
+            let pevals: Vec<F> = peval_shares
+                .into_iter()
+                .flat_map(|x| pp.unpack(x))
+                .collect();
+
+            log::info!("Party {} about to validate", net.party_id());
+            if net.is_king() {
+                assert_eq!(should_be_output, pevals);
+            }
+        }
+
+        log::info!("Party {} done", net.party_id());
+    }
+
+    pub fn setup_log() {
+        let _ = SubscriberBuilder::default()
+            .with_env_filter(EnvFilter::from_default_env())
+            .finish()
+            .try_init();
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_dfft() -> Result<(), Box<dyn Error>> {
+        setup_log();
+        test_gadget::simulate_test(
+            N,
+            10,
+            // Give 10 minutes per test
+            std::time::Duration::from_secs(60 * 10),
+            vec![0],
+            async_proto_generator,
+        )
+        .await
+    }
+
+    fn async_proto_generator(
+        mut params: TestAsyncProtocolParameters<TestBundle>,
+    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+        Box::pin(async move {
+            params
+                .start_rx
+                .take()
+                .expect("Already started")
+                .await
+                .expect("Failed to start");
+            let on_end_tx = params.test_bundle.count_finished_tx.clone();
+            // We need to create a network that implements MpcNet. This means we need multiplexing.
+            // To multiplex the inbound stream from the JobManager, we need to take the receive handle
+            // and spawn a task that will receive messages and send them to the appropriate channel.
+            // For sending messages, the payload we send to the JobManager needs to be multiplexed with
+            // stream IDs.
+
+            let mut txs = HashMap::new();
+            let mut rxs = HashMap::new();
+            for peer_id in 0..params.test_bundle.n_peers {
+                // Create 3 multiplexed channels
+                let mut txs_for_this_peer = vec![];
+                let mut rxs_for_this_peer = vec![];
+                for _ in 0..3 {
+                    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+                    txs_for_this_peer.push(tx);
+                    rxs_for_this_peer.push(Mutex::new(rx));
+                }
+
+                txs.insert(peer_id as u32, txs_for_this_peer);
+                rxs.insert(peer_id as u32, rxs_for_this_peer);
+            }
+
+            let network = ZkNetworkOverGadgetNetwork {
+                gadget_network: params.test_bundle.network,
+                rxs,
+                n_peers: params.test_bundle.n_peers,
+                party_id: params.test_bundle.party_id,
+                associated_block_id: params.associated_block_id,
+                associated_session_id: params.associated_session_id,
+                associated_ssid: params.associated_ssid,
+                associated_task_id: params.associated_task_id,
+            };
+
+            tokio::task::spawn(async move {
+                while let Some(message) = params.protocol_message_rx.recv().await {
+                    println!(
+                        "Message RECV ({} -> {}): {}",
+                        message.from,
+                        message.to.expect("Should exist"),
+                        message.payload.debug()
+                    );
+                    let deserialized: MpcNetMessage =
+                        bincode2::deserialize(&message.payload).expect("Failed to deser message");
+                    assert_ne!(deserialized.source, network.party_id);
+                    assert_eq!(deserialized.source, message.from);
+                    assert_eq!(network.party_id, message.to.expect("Should exist"));
+
+                    let txs = &txs[&deserialized.source];
+                    let tx = &txs[deserialized.sid as usize];
+                    tx.send(deserialized).expect("Failed to send message");
+                }
+            });
+
+            let pp = PackedSharingParams::<Fr>::new(L);
+            let dom = Radix2EvaluationDomain::<Fr>::new(M).unwrap();
+            d_fft_test::<Fr, _>(&pp, &dom, &network).await;
+            on_end_tx.send(()).expect("Failed to send on_end signal");
+        })
+    }
+
+    struct ZkNetworkOverGadgetNetwork {
+        gadget_network: InMemoryNetwork,
+        rxs: HashMap<u32, Vec<Mutex<tokio::sync::mpsc::UnboundedReceiver<MpcNetMessage>>>>,
+        n_peers: usize,
+        party_id: u32,
+        associated_block_id: u64,
+        associated_session_id: u64,
+        associated_ssid: u16,
+        associated_task_id: [u8; 8],
+    }
+
+    #[derive(Serialize, Deserialize)]
+    struct MpcNetMessage {
+        sid: MultiplexedStreamID,
+        payload: bytes::Bytes,
+        source: u32,
+    }
+
+    #[async_trait]
+    impl MpcNet for ZkNetworkOverGadgetNetwork {
+        fn n_parties(&self) -> usize {
+            self.n_peers
+        }
+
+        fn party_id(&self) -> u32 {
+            self.party_id
+        }
+
+        fn is_init(&self) -> bool {
+            true
+        }
+
+        async fn recv_from(&self, id: u32, sid: MultiplexedStreamID) -> Result<Bytes, MpcNetError> {
+            Ok(recv_bytes(sid, id, self).await.payload)
+        }
+
+        async fn send_to(
+            &self,
+            id: u32,
+            bytes: Bytes,
+            sid: MultiplexedStreamID,
+        ) -> Result<(), MpcNetError> {
+            send_bytes(sid, bytes, self, id).await;
+            Ok(())
+        }
+    }
+
+    async fn send_bytes(
+        sid: MultiplexedStreamID,
+        payload: Bytes,
+        network: &ZkNetworkOverGadgetNetwork,
+        to: u32,
+    ) {
+        let mpc_net_payload = MpcNetMessage {
+            sid,
+            payload,
+            source: network.party_id,
+        };
+
+        let serialized =
+            bincode2::serialize(&mpc_net_payload).expect("Failed to serialize message");
+        println!(
+            "Message SEND ({} -> {}): {}",
+            network.party_id,
+            to,
+            serialized.debug()
+        );
+
+        let gadget_protocol_message = TestProtocolMessage {
+            payload: serialized,
+            from: network.party_id,
+            to: Some(to),
+            associated_block_id: network.associated_block_id,
+            associated_session_id: network.associated_session_id,
+            associated_ssid: network.associated_ssid,
+            associated_task_id: network.associated_task_id,
+        };
+
+        assert_ne!(to, network.party_id, "Cannot send to self");
+        network
+            .gadget_network
+            .send_to(gadget_protocol_message, to)
+            .expect("Failed to send");
+    }
+
+    async fn recv_bytes(
+        sid: MultiplexedStreamID,
+        from: u32,
+        network: &ZkNetworkOverGadgetNetwork,
+    ) -> MpcNetMessage {
+        let rx = &network.rxs.get(&from).expect("Should exist")[sid as usize];
+        rx.lock()
+            .await
+            .recv()
+            .await
+            .expect("Failed to receive bytes")
+    }
+
+    pub trait Debuggable: AsRef<[u8]> {
+        fn debug(&self) -> String {
+            let hash = md5::compute(self.as_ref()).0;
+            hex::encode(hash)
+        }
+    }
+
+    impl<T: AsRef<[u8]>> Debuggable for T {}
+}
diff --git a/test-gadget/tests/dmsm.rs b/test-gadget/tests/dmsm.rs
new file mode 100644
index 000000000..bbc361f8b
--- /dev/null
+++ b/test-gadget/tests/dmsm.rs
@@ -0,0 +1,275 @@
+#[cfg(test)]
+mod tests {
+    use dist_primitives::dmsm::d_msm;
+    use dist_primitives::dmsm::packexp_from_public;
+    use gadget_core::job_manager::SendFuture;
+    use std::collections::HashMap;
+    use std::error::Error;
+    use std::pin::Pin;
+    use test_gadget::work_manager::TestAsyncProtocolParameters;
+    use test_gadget::TestBundle;
+    use tracing_subscriber::fmt::SubscriberBuilder;
+    use tracing_subscriber::util::SubscriberInitExt;
+    use tracing_subscriber::EnvFilter;
+
+    use ark_bls12_377::Fr;
+    use ark_ec::CurveGroup;
+    use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
+    use ark_std::UniformRand;
+    use async_trait::async_trait;
+    use bytes::Bytes;
+    use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
+    use secret_sharing::pss::PackedSharingParams;
+    use serde::{Deserialize, Serialize};
+    use test_gadget::message::TestProtocolMessage;
+    use test_gadget::test_network::InMemoryNetwork;
+    use tokio::sync::Mutex;
+
+    const L: usize = 2;
+    const N: usize = L * 4;
+    const M: usize = L * 1024;
+
+    pub async fn d_msm_test<G: CurveGroup, Net: MpcNet + 'static>(
+        pp: PackedSharingParams<G::ScalarField>,
+        dom: Radix2EvaluationDomain<G::ScalarField>,
+        net: Net,
+    ) {
+        let (x_share_aff, y_share, pp, net, should_be_output) =
+            tokio::task::spawn_blocking(move || {
+                // let m = pp.l*4;
+                let mbyl: usize = dom.size() / pp.l;
+                log::info!(
+                    "m: {}, mbyl: {}, party_id: {}",
+                    dom.size(),
+                    mbyl,
+                    net.party_id()
+                );
+
+                let rng = &mut ark_std::test_rng();
+
+                let mut y_pub: Vec<G::ScalarField> = Vec::new();
+                let mut x_pub: Vec<G> = Vec::new();
+
+                for _ in 0..dom.size() {
+                    y_pub.push(G::ScalarField::rand(rng));
+                    x_pub.push(G::rand(rng));
+                }
+
+                log::info!("About to begin packexp_from_public");
+
+                let x_share: Vec<G> = x_pub
+                    .chunks(pp.l)
+                    .map(|s| packexp_from_public(s, &pp)[net.party_id() as usize])
+                    .collect();
+
+                log::info!("About to begin pack_from_public");
+                let y_share: Vec<G::ScalarField> = y_pub
+                    .chunks(pp.l)
+                    .map(|s| pp.pack_from_public(s.to_vec())[net.party_id() as usize])
+                    .collect();
+
+                let x_pub_aff: Vec<G::Affine> = x_pub.iter().map(|s| (*s).into()).collect();
+                let x_share_aff: Vec<G::Affine> = x_share.iter().map(|s| (*s).into()).collect();
+
+                // Will be comparing against this in the end
+                log::info!("About to begin G::msm");
+                let should_be_output = G::msm(x_pub_aff.as_slice(), y_pub.as_slice()).unwrap();
+                (x_share_aff, y_share, pp, net, should_be_output)
+            })
+            .await
+            .expect("Failed to spawn blocking task");
+
+        log::info!("About to begin d_msm");
+        let output = d_msm::<G, Net>(&x_share_aff, &y_share, &pp, &net, MultiplexedStreamID::Zero)
+            .await
+            .unwrap();
+
+        if net.is_king() {
+            log::info!("About to king assert");
+            assert_eq!(should_be_output, output);
+        }
+
+        log::info!("Party {} done", net.party_id());
+    }
+
+    pub fn setup_log() {
+        let _ = SubscriberBuilder::default()
+            .with_env_filter(EnvFilter::from_default_env())
+            .finish()
+            .try_init();
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_dmsm() -> Result<(), Box<dyn Error>> {
+        setup_log();
+        test_gadget::simulate_test(
+            N,
+            10,
+            // Give 10 minutes per test
+            std::time::Duration::from_secs(60 * 10),
+            vec![0],
+            async_proto_generator,
+        )
+        .await
+    }
+
+    fn async_proto_generator(
+        mut params: TestAsyncProtocolParameters<TestBundle>,
+    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+        Box::pin(async move {
+            params
+                .start_rx
+                .take()
+                .expect("Already started")
+                .await
+                .expect("Failed to start");
+            let on_end_tx = params.test_bundle.count_finished_tx.clone();
+            // We need to create a network that implements MpcNet. This means we need multiplexing.
+            // To multiplex the inbound stream from the JobManager, we need to take the receive handle
+            // and spawn a task that will receive messages and send them to the appropriate channel.
+            // For sending messages, the payload we send to the JobManager needs to be multiplexed with
+            // stream IDs.
+
+            let mut txs = HashMap::new();
+            let mut rxs = HashMap::new();
+            for peer_id in 0..params.test_bundle.n_peers {
+                // Create 3 multiplexed channels
+                let mut txs_for_this_peer = vec![];
+                let mut rxs_for_this_peer = vec![];
+                for _ in 0..3 {
+                    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+                    txs_for_this_peer.push(tx);
+                    rxs_for_this_peer.push(Mutex::new(rx));
+                }
+
+                txs.insert(peer_id as u32, txs_for_this_peer);
+                rxs.insert(peer_id as u32, rxs_for_this_peer);
+            }
+
+            let network = ZkNetworkOverGadgetNetwork {
+                gadget_network: params.test_bundle.network,
+                rxs,
+                n_peers: params.test_bundle.n_peers,
+                party_id: params.test_bundle.party_id,
+                associated_block_id: params.associated_block_id,
+                associated_session_id: params.associated_session_id,
+                associated_ssid: params.associated_ssid,
+                associated_task_id: params.associated_task_id,
+            };
+
+            tokio::task::spawn(async move {
+                while let Some(message) = params.protocol_message_rx.recv().await {
+                    let deserialized: MpcNetMessage =
+                        bincode2::deserialize(&message.payload).expect("Failed to deser message");
+                    let txs = &txs[&deserialized.source];
+                    let tx = &txs[deserialized.sid as usize];
+                    tx.send(deserialized).expect("Failed to send message");
+                }
+            });
+
+            let pp = PackedSharingParams::<Fr>::new(L);
+            let dom = Radix2EvaluationDomain::<Fr>::new(M).unwrap();
+            d_msm_test::<ark_bls12_377::G1Projective, _>(pp, dom, network).await;
+            on_end_tx.send(()).expect("Failed to send on_end signal");
+        })
+    }
+
+    struct ZkNetworkOverGadgetNetwork {
+        gadget_network: InMemoryNetwork,
+        rxs: HashMap<u32, Vec<Mutex<tokio::sync::mpsc::UnboundedReceiver<MpcNetMessage>>>>,
+        n_peers: usize,
+        party_id: u32,
+        associated_block_id: u64,
+        associated_session_id: u64,
+        associated_ssid: u16,
+        associated_task_id: [u8; 8],
+    }
+
+    #[derive(Serialize, Deserialize)]
+    struct MpcNetMessage {
+        sid: MultiplexedStreamID,
+        payload: Bytes,
+        source: u32,
+    }
+
+    #[async_trait]
+    impl MpcNet for ZkNetworkOverGadgetNetwork {
+        fn n_parties(&self) -> usize {
+            self.n_peers
+        }
+
+        fn party_id(&self) -> u32 {
+            self.party_id
+        }
+
+        fn is_init(&self) -> bool {
+            true
+        }
+
+        async fn recv_from(&self, id: u32, sid: MultiplexedStreamID) -> Result<Bytes, MpcNetError> {
+            Ok(recv_bytes(sid, id, self).await.payload)
+        }
+
+        async fn send_to(
+            &self,
+            id: u32,
+            bytes: Bytes,
+            sid: MultiplexedStreamID,
+        ) -> Result<(), MpcNetError> {
+            send_bytes(sid, bytes, self, Some(id)).await;
+            Ok(())
+        }
+    }
+
+    async fn send_bytes(
+        sid: MultiplexedStreamID,
+        payload: Bytes,
+        network: &ZkNetworkOverGadgetNetwork,
+        to: Option<u32>,
+    ) {
+        let mpc_net_payload = MpcNetMessage {
+            sid,
+            payload,
+            source: network.party_id,
+        };
+
+        let serialized =
+            bincode2::serialize(&mpc_net_payload).expect("Failed to serialize message");
+
+        let gadget_protocol_message = TestProtocolMessage {
+            payload: serialized,
+            from: network.party_id,
+            to,
+            associated_block_id: network.associated_block_id,
+            associated_session_id: network.associated_session_id,
+            associated_ssid: network.associated_ssid,
+            associated_task_id: network.associated_task_id,
+        };
+
+        if let Some(to) = to {
+            assert_ne!(to, network.party_id, "Cannot send to self");
+            network
+                .gadget_network
+                .send_to(gadget_protocol_message, to)
+                .expect("Failed to send");
+        } else {
+            network
+                .gadget_network
+                .broadcast(network.party_id, gadget_protocol_message)
+                .expect("Failed to broadcast");
+        }
+    }
+
+    async fn recv_bytes(
+        sid: MultiplexedStreamID,
+        from: u32,
+        network: &ZkNetworkOverGadgetNetwork,
+    ) -> MpcNetMessage {
+        let rx = &network.rxs.get(&from).expect("Should exist")[sid as usize];
+        rx.lock()
+            .await
+            .recv()
+            .await
+            .expect("Failed to receive bytes")
+    }
+}
diff --git a/test-gadget/tests/zk_network.rs b/test-gadget/tests/zk_network.rs
new file mode 100644
index 000000000..273f88d77
--- /dev/null
+++ b/test-gadget/tests/zk_network.rs
@@ -0,0 +1,249 @@
+#[cfg(test)]
+mod tests {
+    use gadget_core::job_manager::SendFuture;
+    use std::collections::HashMap;
+    use std::error::Error;
+    use std::pin::Pin;
+    use test_gadget::work_manager::TestAsyncProtocolParameters;
+    use test_gadget::TestBundle;
+    use tracing_subscriber::fmt::SubscriberBuilder;
+    use tracing_subscriber::util::SubscriberInitExt;
+    use tracing_subscriber::EnvFilter;
+
+    use async_trait::async_trait;
+    use bytes::Bytes;
+    use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
+    use serde::{Deserialize, Serialize};
+    use test_gadget::message::{TestProtocolMessage, UserID};
+    use test_gadget::test_network::InMemoryNetwork;
+    use tokio::sync::Mutex;
+
+    pub fn setup_log() {
+        let _ = SubscriberBuilder::default()
+            .with_env_filter(EnvFilter::from_default_env())
+            .finish()
+            .try_init();
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_zk_network() -> Result<(), Box<dyn Error>> {
+        setup_log();
+        test_gadget::simulate_test(
+            5,
+            10,
+            std::time::Duration::from_secs(60 * 10),
+            vec![0],
+            async_proto_generator,
+        )
+        .await
+    }
+
+    fn async_proto_generator(
+        mut params: TestAsyncProtocolParameters<TestBundle>,
+    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+        Box::pin(async move {
+            params
+                .start_rx
+                .take()
+                .expect("Already started")
+                .await
+                .expect("Failed to start");
+            let on_end_tx = params.test_bundle.count_finished_tx.clone();
+            // We need to create a network that implements MpcNet. This means we need multiplexing.
+            // To multiplex the inbound stream from the JobManager, we need to take the receive handle
+            // and spawn a task that will receive messages and send them to the appropriate channel.
+            // For sending messages, the payload we send to the JobManager needs to be multiplexed with
+            // stream IDs.
+
+            let mut txs = HashMap::new();
+            let mut rxs = HashMap::new();
+            for peer_id in 0..params.test_bundle.n_peers {
+                // Create 3 multiplexed channels
+                let mut txs_for_this_peer = vec![];
+                let mut rxs_for_this_peer = vec![];
+                for _ in 0..3 {
+                    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+                    txs_for_this_peer.push(tx);
+                    rxs_for_this_peer.push(Mutex::new(rx));
+                }
+
+                txs.insert(peer_id as u32, txs_for_this_peer);
+                rxs.insert(peer_id as u32, rxs_for_this_peer);
+            }
+
+            let network = ZkNetworkOverGadgetNetwork {
+                gadget_network: params.test_bundle.network,
+                rxs,
+                n_peers: params.test_bundle.n_peers,
+                party_id: params.test_bundle.party_id,
+                associated_block_id: params.associated_block_id,
+                associated_session_id: params.associated_session_id,
+                associated_ssid: params.associated_ssid,
+                associated_task_id: params.associated_task_id,
+            };
+
+            tokio::task::spawn(async move {
+                while let Some(message) = params.protocol_message_rx.recv().await {
+                    let deserialized: MpcNetMessage =
+                        bincode2::deserialize(&message.payload).expect("Failed to deser message");
+                    let txs = &txs[&deserialized.source];
+                    let tx = &txs[deserialized.sid as usize];
+                    tx.send(deserialized).expect("Failed to send message");
+                }
+            });
+
+            let expected_sum = (0..params.test_bundle.n_peers)
+                .map(|i| i as UserID)
+                .sum::<UserID>();
+
+            for sid in [
+                MultiplexedStreamID::Zero,
+                MultiplexedStreamID::One,
+                MultiplexedStreamID::Two,
+            ] {
+                let message = bincode2::serialize(&params.test_bundle.party_id)
+                    .expect("Failed to serialize message");
+                let king_response = if let Some(messages) = network
+                    .client_send_or_king_receive(&message, sid)
+                    .await
+                    .expect("Failed to send")
+                {
+                    assert_eq!(messages.len(), params.test_bundle.n_peers);
+                    let mut sum = 0;
+                    for message in messages.into_iter() {
+                        let peer_id: UserID =
+                            bincode2::deserialize(&message).expect("Failed to deserialize message");
+                        sum += peer_id
+                    }
+                    assert_eq!(
+                        sum, expected_sum,
+                        "Sum of peer IDs should be equal to sum of all peer IDs"
+                    );
+
+                    let sum = Bytes::from(
+                        bincode2::serialize(&sum).expect("Failed to serialize message"),
+                    );
+                    Some(
+                        (0..params.test_bundle.n_peers)
+                            .map(|_| sum.clone())
+                            .collect::<Vec<_>>(),
+                    )
+                } else {
+                    None
+                };
+
+                let sum = network
+                    .client_receive_or_king_send(king_response, sid)
+                    .await
+                    .expect("Failed to receive");
+                let sum: UserID =
+                    bincode2::deserialize(&sum).expect("Failed to deserialize message");
+                assert_eq!(
+                    sum, expected_sum,
+                    "Sum of peer IDs should be equal to sum of all peer IDs"
+                );
+            }
+
+            on_end_tx.send(()).expect("Failed to send on_end signal");
+        })
+    }
+
+    struct ZkNetworkOverGadgetNetwork {
+        gadget_network: InMemoryNetwork,
+        rxs: HashMap<u32, Vec<Mutex<tokio::sync::mpsc::UnboundedReceiver<MpcNetMessage>>>>,
+        n_peers: usize,
+        party_id: u32,
+        associated_block_id: u64,
+        associated_session_id: u64,
+        associated_ssid: u16,
+        associated_task_id: [u8; 8],
+    }
+
+    #[derive(Serialize, Deserialize)]
+    struct MpcNetMessage {
+        sid: MultiplexedStreamID,
+        payload: Bytes,
+        source: u32,
+    }
+
+    #[async_trait]
+    impl MpcNet for ZkNetworkOverGadgetNetwork {
+        fn n_parties(&self) -> usize {
+            self.n_peers
+        }
+
+        fn party_id(&self) -> u32 {
+            self.party_id
+        }
+
+        fn is_init(&self) -> bool {
+            true
+        }
+
+        async fn recv_from(&self, id: u32, sid: MultiplexedStreamID) -> Result<Bytes, MpcNetError> {
+            Ok(recv_bytes(sid, id, self).await.payload)
+        }
+
+        async fn send_to(
+            &self,
+            id: u32,
+            bytes: Bytes,
+            sid: MultiplexedStreamID,
+        ) -> Result<(), MpcNetError> {
+            send_bytes(sid, bytes, self, Some(id)).await;
+            Ok(())
+        }
+    }
+
+    async fn send_bytes(
+        sid: MultiplexedStreamID,
+        payload: Bytes,
+        network: &ZkNetworkOverGadgetNetwork,
+        to: Option<u32>,
+    ) {
+        let mpc_net_payload = MpcNetMessage {
+            sid,
+            payload,
+            source: network.party_id,
+        };
+
+        let serialized =
+            bincode2::serialize(&mpc_net_payload).expect("Failed to serialize message");
+
+        let gadget_protocol_message = TestProtocolMessage {
+            payload: serialized,
+            from: network.party_id,
+            to,
+            associated_block_id: network.associated_block_id,
+            associated_session_id: network.associated_session_id,
+            associated_ssid: network.associated_ssid,
+            associated_task_id: network.associated_task_id,
+        };
+
+        if let Some(to) = to {
+            assert_ne!(to, network.party_id, "Cannot send to self");
+            network
+                .gadget_network
+                .send_to(gadget_protocol_message, to)
+                .expect("Failed to send");
+        } else {
+            network
+                .gadget_network
+                .broadcast(network.party_id, gadget_protocol_message)
+                .expect("Failed to broadcast");
+        }
+    }
+
+    async fn recv_bytes(
+        sid: MultiplexedStreamID,
+        from: u32,
+        network: &ZkNetworkOverGadgetNetwork,
+    ) -> MpcNetMessage {
+        let rx = &network.rxs.get(&from).expect("Should exist")[sid as usize];
+        rx.lock()
+            .await
+            .recv()
+            .await
+            .expect("Failed to receive bytes")
+    }
+}