diff --git a/Cargo.toml b/Cargo.toml
index af374bc..bbd0c12 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "wsts"
-version = "10.0.0"
+version = "11.0.0"
 edition = "2021"
 authors = ["Joey Yandle <xoloki@gmail.com>"]
 license = "Apache-2.0"
diff --git a/src/common.rs b/src/common.rs
index 89e3570..fad2216 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -321,7 +321,7 @@ pub mod test_helpers {
 
 #[cfg(test)]
 /// Test module for common functionality
-pub mod test {
+mod test {
     use super::*;
     use crate::util::create_rng;
 
diff --git a/src/net.rs b/src/net.rs
index 2d305ff..464a39b 100644
--- a/src/net.rs
+++ b/src/net.rs
@@ -168,6 +168,9 @@ pub struct DkgPrivateBegin {
     pub signer_ids: Vec<u32>,
     /// Key IDs who responded in time for this DKG round
     pub key_ids: Vec<u32>,
+    /// Include DkgPublicShares to avoid p2p related message delivery
+    /// order issues when signers communicate directly with each other
+    pub dkg_public_shares: HashMap<u32, DkgPublicShares>,
 }
 
 impl Signable for DkgPrivateBegin {
@@ -179,6 +182,9 @@ impl Signable for DkgPrivateBegin {
         }
         for signer_id in &self.signer_ids {
             hasher.update(signer_id.to_be_bytes());
+            if let Some(shares) = self.dkg_public_shares.get(signer_id) {
+                shares.hash(hasher);
+            }
         }
     }
 }
@@ -228,6 +234,9 @@ pub struct DkgEndBegin {
     pub signer_ids: Vec<u32>,
     /// Key IDs who responded in time for this DKG round
     pub key_ids: Vec<u32>,
+    /// Include DkgPrivateShares to avoid p2p related message delivery
+    /// order issues when signers communicate directly with each other
+    pub dkg_private_shares: HashMap<u32, DkgPrivateShares>,
 }
 
 impl Signable for DkgEndBegin {
@@ -239,6 +248,9 @@ impl Signable for DkgEndBegin {
         }
         for signer_id in &self.signer_ids {
             hasher.update(signer_id.to_be_bytes());
+            if let Some(shares) = self.dkg_private_shares.get(signer_id) {
+                shares.hash(hasher);
+            }
         }
     }
 }
@@ -660,6 +672,7 @@ mod test {
             dkg_id: 0,
             key_ids: Default::default(),
             signer_ids: Default::default(),
+            dkg_public_shares: Default::default(),
         };
         let msg = Message::DkgBegin(dkg_begin.clone());
         let coordinator_packet_dkg_begin = Packet {
diff --git a/src/state_machine/coordinator/fire.rs b/src/state_machine/coordinator/fire.rs
index 0446096..3293e88 100644
--- a/src/state_machine/coordinator/fire.rs
+++ b/src/state_machine/coordinator/fire.rs
@@ -424,11 +424,19 @@ impl<Aggregator: AggregatorTrait> Coordinator<Aggregator> {
             .keys()
             .flat_map(|signer_id| self.config.signer_key_ids[signer_id].clone())
             .collect::<Vec<u32>>();
-
+        let dkg_public_shares = if self.config.embed_public_private_shares {
+            self.dkg_public_shares
+                .iter()
+                .map(|(id, share)| (*id, share.clone()))
+                .collect()
+        } else {
+            Default::default()
+        };
         let dkg_begin = DkgPrivateBegin {
             dkg_id: self.current_dkg_id,
             key_ids: active_key_ids,
             signer_ids: self.dkg_public_shares.keys().cloned().collect(),
+            dkg_public_shares,
         };
         let dkg_private_begin_msg = Packet {
             sig: dkg_begin
@@ -458,11 +466,20 @@ impl<Aggregator: AggregatorTrait> Coordinator<Aggregator> {
             .keys()
             .flat_map(|signer_id| self.config.signer_key_ids[signer_id].clone())
             .collect::<Vec<u32>>();
+        let dkg_private_shares = if self.config.embed_public_private_shares {
+            self.dkg_private_shares
+                .iter()
+                .map(|(id, share)| (*id, share.clone()))
+                .collect()
+        } else {
+            Default::default()
+        };
 
         let dkg_end_begin = DkgEndBegin {
             dkg_id: self.current_dkg_id,
             key_ids: active_key_ids,
             signer_ids: self.dkg_private_shares.keys().cloned().collect(),
+            dkg_private_shares,
         };
         let dkg_end_begin_msg = Packet {
             sig: dkg_end_begin
@@ -1270,7 +1287,7 @@ impl<Aggregator: AggregatorTrait> CoordinatorTrait for Coordinator<Aggregator> {
 
 #[cfg(test)]
 /// Test module for coordinator functionality
-pub mod test {
+mod test {
     use crate::{
         curve::{point::Point, scalar::Scalar},
         net::{
@@ -1524,17 +1541,27 @@ pub mod test {
 
     #[test]
     fn missing_public_keys_dkg_v1() {
-        missing_public_keys_dkg::<v1::Aggregator, v1::Signer>(10, 1);
+        missing_public_keys_dkg::<v1::Aggregator, v1::Signer>(10, 1, false);
+    }
+    #[test]
+    fn missing_public_keys_dkg_v1_embed() {
+        missing_public_keys_dkg::<v1::Aggregator, v1::Signer>(10, 1, true);
     }
 
     #[test]
     fn missing_public_keys_dkg_v2() {
-        missing_public_keys_dkg::<v2::Aggregator, v2::Signer>(10, 1);
+        missing_public_keys_dkg::<v2::Aggregator, v2::Signer>(10, 1, false);
+    }
+
+    #[test]
+    fn missing_public_keys_dkg_v2_embed() {
+        missing_public_keys_dkg::<v2::Aggregator, v2::Signer>(10, 1, true);
     }
 
     fn missing_public_keys_dkg<Aggregator: AggregatorTrait, SignerType: SignerTrait>(
         num_signers: u32,
         keys_per_signer: u32,
+        embed_public_private_shares: bool,
     ) -> (Vec<FireCoordinator<Aggregator>>, Vec<Signer<SignerType>>) {
         let timeout = Duration::from_millis(1024);
         let expire = Duration::from_millis(1280);
@@ -1547,6 +1574,7 @@ pub mod test {
                 Some(timeout),
                 Some(timeout),
                 Some(timeout),
+                embed_public_private_shares,
             );
 
         // Start a DKG round where we will not allow all signers to recv DkgBegin, so they will not respond with DkgPublicShares
@@ -1605,17 +1633,28 @@ pub mod test {
 
     #[test]
     fn minimum_signers_dkg_v1() {
-        minimum_signers_dkg::<v1::Aggregator, v1::Signer>(10, 2);
+        minimum_signers_dkg::<v1::Aggregator, v1::Signer>(10, 2, false);
     }
 
     #[test]
     fn minimum_signers_dkg_v2() {
-        minimum_signers_dkg::<v2::Aggregator, v2::Signer>(10, 2);
+        minimum_signers_dkg::<v2::Aggregator, v2::Signer>(10, 2, false);
+    }
+
+    #[test]
+    fn minimum_signers_dkg_v1_embed() {
+        minimum_signers_dkg::<v1::Aggregator, v1::Signer>(10, 2, true);
+    }
+
+    #[test]
+    fn minimum_signers_dkg_v2_embed() {
+        minimum_signers_dkg::<v2::Aggregator, v2::Signer>(10, 2, true);
     }
 
     fn minimum_signers_dkg<Aggregator: AggregatorTrait, SignerType: SignerTrait>(
         num_signers: u32,
         keys_per_signer: u32,
+        embed_public_private_shares: bool,
     ) -> (Vec<FireCoordinator<Aggregator>>, Vec<Signer<SignerType>>) {
         let timeout = Duration::from_millis(1024);
         let expire = Duration::from_millis(1280);
@@ -1628,6 +1667,7 @@ pub mod test {
                 Some(timeout),
                 Some(timeout),
                 Some(timeout),
+                embed_public_private_shares,
             );
 
         // Start a DKG round where we will not allow all signers to recv DkgBegin, so they will not respond with DkgPublicShares
@@ -1760,15 +1800,27 @@ pub mod test {
 
     #[test]
     fn insufficient_signers_dkg_v1() {
-        insufficient_signers_dkg::<v1::Aggregator, v1::Signer>();
+        insufficient_signers_dkg::<v1::Aggregator, v1::Signer>(false);
     }
 
     #[test]
     fn insufficient_signers_dkg_v2() {
-        insufficient_signers_dkg::<v2::Aggregator, v2::Signer>();
+        insufficient_signers_dkg::<v2::Aggregator, v2::Signer>(false);
+    }
+
+    #[test]
+    fn insufficient_signers_dkg_v1_embed() {
+        insufficient_signers_dkg::<v1::Aggregator, v1::Signer>(true);
+    }
+
+    #[test]
+    fn insufficient_signers_dkg_v2_embed() {
+        insufficient_signers_dkg::<v2::Aggregator, v2::Signer>(true);
     }
 
-    fn insufficient_signers_dkg<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
+    fn insufficient_signers_dkg<Aggregator: AggregatorTrait, Signer: SignerTrait>(
+        embed_public_private_shares: bool,
+    ) {
         let timeout = Duration::from_millis(1024);
         let expire = Duration::from_millis(1280);
         let num_signers = 10;
@@ -1781,6 +1833,7 @@ pub mod test {
             Some(timeout),
             Some(timeout),
             Some(timeout),
+            embed_public_private_shares,
         );
 
         // Start a DKG round where we will not allow all signers to recv DkgBegin, so they will not respond with DkgPublicShares
@@ -2209,20 +2262,35 @@ pub mod test {
 
     #[test]
     fn minimum_signers_sign_v1() {
-        minimum_signers_sign::<v1::Aggregator, v1::Signer>();
+        minimum_signers_sign::<v1::Aggregator, v1::Signer>(false);
     }
 
     #[test]
     fn minimum_signers_sign_v2() {
-        minimum_signers_sign::<v2::Aggregator, v2::Signer>();
+        minimum_signers_sign::<v2::Aggregator, v2::Signer>(false);
     }
 
-    fn minimum_signers_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
+    #[test]
+    fn minimum_signers_sign_v1_embed() {
+        minimum_signers_sign::<v1::Aggregator, v1::Signer>(true);
+    }
+
+    #[test]
+    fn minimum_signers_sign_v2_embed() {
+        minimum_signers_sign::<v2::Aggregator, v2::Signer>(true);
+    }
+
+    fn minimum_signers_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>(
+        embed_public_private_shares: bool,
+    ) {
         let num_signers = 10;
         let keys_per_signer = 2;
 
-        let (mut coordinators, mut signers) =
-            minimum_signers_dkg::<Aggregator, Signer>(num_signers, keys_per_signer);
+        let (mut coordinators, mut signers) = minimum_signers_dkg::<Aggregator, Signer>(
+            num_signers,
+            keys_per_signer,
+            embed_public_private_shares,
+        );
         let config = coordinators.first().unwrap().get_config();
 
         // Figure out how many signers we can remove and still be above the threshold
@@ -2293,20 +2361,35 @@ pub mod test {
 
     #[test]
     fn missing_public_keys_sign_v1() {
-        missing_public_keys_sign::<v1::Aggregator, v1::Signer>();
+        missing_public_keys_sign::<v1::Aggregator, v1::Signer>(false);
+    }
+
+    #[test]
+    fn missing_public_keys_sign_v2() {
+        missing_public_keys_sign::<v2::Aggregator, v2::Signer>(false);
     }
 
     #[test]
-    fn minimum_missing_public_keys_sign_v2() {
-        missing_public_keys_sign::<v2::Aggregator, v2::Signer>();
+    fn missing_public_keys_sign_v1_embed() {
+        missing_public_keys_sign::<v1::Aggregator, v1::Signer>(true);
     }
 
-    fn missing_public_keys_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
+    #[test]
+    fn missing_public_keys_sign_v2_embed() {
+        missing_public_keys_sign::<v2::Aggregator, v2::Signer>(true);
+    }
+
+    fn missing_public_keys_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>(
+        embed_public_private_shares: bool,
+    ) {
         let num_signers = 10;
         let keys_per_signer = 2;
 
-        let (mut coordinators, mut signers) =
-            minimum_signers_dkg::<Aggregator, Signer>(num_signers, keys_per_signer);
+        let (mut coordinators, mut signers) = minimum_signers_dkg::<Aggregator, Signer>(
+            num_signers,
+            keys_per_signer,
+            embed_public_private_shares,
+        );
 
         // Let us also remove that signers public key from the config including all of its key ids
         let mut removed_signer = signers.pop().expect("Failed to pop signer");
@@ -2380,15 +2463,27 @@ pub mod test {
 
     #[test]
     fn insufficient_signers_sign_v1() {
-        insufficient_signers_sign::<v1::Aggregator, v1::Signer>();
+        insufficient_signers_sign::<v1::Aggregator, v1::Signer>(false);
     }
 
     #[test]
     fn insufficient_signers_sign_v2() {
-        insufficient_signers_sign::<v2::Aggregator, v2::Signer>();
+        insufficient_signers_sign::<v2::Aggregator, v2::Signer>(false);
+    }
+
+    #[test]
+    fn insufficient_signers_sign_v1_embed() {
+        insufficient_signers_sign::<v1::Aggregator, v1::Signer>(true);
+    }
+
+    #[test]
+    fn insufficient_signers_sign_v2_embed() {
+        insufficient_signers_sign::<v2::Aggregator, v2::Signer>(true);
     }
 
-    fn insufficient_signers_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>() {
+    fn insufficient_signers_sign<Aggregator: AggregatorTrait, Signer: SignerTrait>(
+        embed_public_private_shares: bool,
+    ) {
         let num_signers = 5;
         let keys_per_signer = 2;
         let (mut coordinators, mut signers) =
@@ -2400,6 +2495,7 @@ pub mod test {
                 None,
                 Some(Duration::from_millis(128)),
                 Some(Duration::from_millis(128)),
+                embed_public_private_shares,
             );
         let config = coordinators.first().unwrap().get_config();
 
diff --git a/src/state_machine/coordinator/frost.rs b/src/state_machine/coordinator/frost.rs
index c4b6145..a5831fd 100644
--- a/src/state_machine/coordinator/frost.rs
+++ b/src/state_machine/coordinator/frost.rs
@@ -229,10 +229,18 @@ impl<Aggregator: AggregatorTrait> Coordinator<Aggregator> {
             dkg_id = %self.current_dkg_id,
             "Starting Private Share Distribution"
         );
+        let dkg_public_shares = if self.config.embed_public_private_shares {
+            (0..self.config.num_signers)
+                .map(|id| (id, self.dkg_public_shares[&id].clone()))
+                .collect()
+        } else {
+            Default::default()
+        };
         let dkg_begin = DkgPrivateBegin {
             dkg_id: self.current_dkg_id,
             key_ids: (1..self.config.num_keys + 1).collect(),
             signer_ids: (0..self.config.num_signers).collect(),
+            dkg_public_shares,
         };
         let dkg_private_begin_msg = Packet {
             sig: dkg_begin
@@ -251,10 +259,18 @@ impl<Aggregator: AggregatorTrait> Coordinator<Aggregator> {
             dkg_id = %self.current_dkg_id,
             "Starting DKG End Distribution"
         );
+        let dkg_private_shares = if self.config.embed_public_private_shares {
+            (0..self.config.num_signers)
+                .map(|id| (id, self.dkg_private_shares[&id].clone()))
+                .collect()
+        } else {
+            Default::default()
+        };
         let dkg_begin = DkgEndBegin {
             dkg_id: self.current_dkg_id,
             key_ids: (0..self.config.num_keys).collect(),
             signer_ids: (0..self.config.num_signers).collect(),
+            dkg_private_shares,
         };
         let dkg_end_begin_msg = Packet {
             sig: dkg_begin.sign(&self.config.message_private_key).expect(""),
@@ -787,7 +803,7 @@ impl<Aggregator: AggregatorTrait> CoordinatorTrait for Coordinator<Aggregator> {
 
 #[cfg(test)]
 /// Test module for coordinator functionality
-pub mod test {
+mod test {
     use crate::{
         curve::scalar::Scalar,
         net::{DkgBegin, Message, NonceRequest, Packet, SignatureType},
diff --git a/src/state_machine/coordinator/mod.rs b/src/state_machine/coordinator/mod.rs
index 9170f4d..b2759ef 100644
--- a/src/state_machine/coordinator/mod.rs
+++ b/src/state_machine/coordinator/mod.rs
@@ -124,6 +124,8 @@ pub struct Config {
     pub signer_key_ids: HashMap<u32, HashSet<u32>>,
     /// ECDSA public keys as Point objects indexed by signer_id
     pub signer_public_keys: HashMap<u32, Point>,
+    /// Embed public and private shares into DkgBegin messages
+    pub embed_public_private_shares: bool,
 }
 
 impl Config {
@@ -147,6 +149,7 @@ impl Config {
             sign_timeout: None,
             signer_key_ids: Default::default(),
             signer_public_keys: Default::default(),
+            embed_public_private_shares: false,
         }
     }
 
@@ -165,6 +168,7 @@ impl Config {
         sign_timeout: Option<Duration>,
         signer_key_ids: HashMap<u32, HashSet<u32>>,
         signer_public_keys: HashMap<u32, Point>,
+        embed_public_private_shares: bool,
     ) -> Self {
         Config {
             num_signers,
@@ -179,6 +183,7 @@ impl Config {
             sign_timeout,
             signer_key_ids,
             signer_public_keys,
+            embed_public_private_shares,
         }
     }
 }
@@ -437,9 +442,12 @@ pub mod test {
             None,
             None,
             None,
+            false,
         )
     }
 
+    #[allow(static_mut_refs)]
+    #[allow(clippy::too_many_arguments)]
     pub fn setup_with_timeouts<Coordinator: CoordinatorTrait, SignerType: SignerTrait>(
         num_signers: u32,
         keys_per_signer: u32,
@@ -448,6 +456,7 @@ pub mod test {
         dkg_end_timeout: Option<Duration>,
         nonce_timeout: Option<Duration>,
         sign_timeout: Option<Duration>,
+        embed_public_private_shares: bool,
     ) -> (Vec<Coordinator>, Vec<Signer<SignerType>>) {
         INIT.call_once(|| {
             tracing_subscriber::registry()
@@ -504,6 +513,7 @@ pub mod test {
                     signer_key_ids[&(signer_id as u32)].clone(),
                     *private_key,
                     public_keys.clone(),
+                    embed_public_private_shares,
                     &mut rng,
                 )
             })
@@ -524,6 +534,7 @@ pub mod test {
                     sign_timeout,
                     signer_key_ids_set.clone(),
                     signer_public_keys.clone(),
+                    embed_public_private_shares,
                 );
                 Coordinator::new(config)
             })
diff --git a/src/state_machine/signer/mod.rs b/src/state_machine/signer/mod.rs
index 61cefde..0d413d4 100644
--- a/src/state_machine/signer/mod.rs
+++ b/src/state_machine/signer/mod.rs
@@ -117,6 +117,8 @@ pub struct SavedState {
     pub dkg_private_begin_msg: Option<DkgPrivateBegin>,
     /// the DKG end begin message received in this round
     pub dkg_end_begin_msg: Option<DkgEndBegin>,
+    /// Embed public and private shares into DkgBegin messages
+    pub embed_public_private_shares: bool,
 }
 
 /// A state machine for a signing round
@@ -166,6 +168,8 @@ pub struct Signer<SignerType: SignerTrait> {
     pub dkg_private_begin_msg: Option<DkgPrivateBegin>,
     /// the DKG end begin message received in this round
     pub dkg_end_begin_msg: Option<DkgEndBegin>,
+    /// Embed public and private shares into DkgBegin messages
+    pub embed_public_private_shares: bool,
 }
 
 impl<SignerType: SignerTrait> Signer<SignerType> {
@@ -179,6 +183,7 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
         key_ids: Vec<u32>,
         network_private_key: Scalar,
         public_keys: PublicKeys,
+        embed_public_private_shares: bool,
         rng: &mut R,
     ) -> Self {
         assert!(threshold <= total_keys);
@@ -216,6 +221,7 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
             dkg_private_shares: Default::default(),
             dkg_private_begin_msg: Default::default(),
             dkg_end_begin_msg: Default::default(),
+            embed_public_private_shares,
         }
     }
 
@@ -242,6 +248,7 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
             dkg_private_shares: state.dkg_private_shares.clone(),
             dkg_private_begin_msg: state.dkg_private_begin_msg.clone(),
             dkg_end_begin_msg: state.dkg_end_begin_msg.clone(),
+            embed_public_private_shares: state.embed_public_private_shares,
         }
     }
 
@@ -268,6 +275,7 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
             dkg_private_shares: self.dkg_private_shares.clone(),
             dkg_private_begin_msg: self.dkg_private_begin_msg.clone(),
             dkg_end_begin_msg: self.dkg_end_begin_msg.clone(),
+            embed_public_private_shares: self.embed_public_private_shares,
         }
     }
 
@@ -320,10 +328,20 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
             Message::DkgPrivateBegin(dkg_private_begin) => {
                 self.dkg_private_begin(dkg_private_begin, rng)
             }
-            Message::DkgEndBegin(dkg_end_begin) => self.dkg_end_begin(dkg_end_begin),
-            Message::DkgPublicShares(dkg_public_shares) => self.dkg_public_share(dkg_public_shares),
+            Message::DkgEndBegin(dkg_end_begin) => self.dkg_end_begin(dkg_end_begin, rng),
+            Message::DkgPublicShares(dkg_public_shares) => {
+                if !self.embed_public_private_shares {
+                    self.dkg_public_share(dkg_public_shares)
+                } else {
+                    Ok(vec![])
+                }
+            }
             Message::DkgPrivateShares(dkg_private_shares) => {
-                self.dkg_private_shares(dkg_private_shares, rng)
+                if !self.embed_public_private_shares {
+                    self.dkg_private_shares(dkg_private_shares, rng)
+                } else {
+                    Ok(vec![])
+                }
             }
             Message::SignatureShareRequest(sign_share_request) => {
                 self.sign_share_request(sign_share_request)
@@ -698,6 +716,10 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
             .cloned()
             .collect::<HashSet<u32>>();
 
+        for (_, shares) in &dkg_private_begin.dkg_public_shares {
+            let _ = self.dkg_public_share(shares)?;
+        }
+
         self.dkg_private_begin_msg = Some(dkg_private_begin.clone());
         self.move_to(State::DkgPrivateDistribute)?;
 
@@ -746,11 +768,19 @@ impl<SignerType: SignerTrait> Signer<SignerType> {
     }
 
     /// handle incoming DkgEndBegin
-    pub fn dkg_end_begin(&mut self, dkg_end_begin: &DkgEndBegin) -> Result<Vec<Message>, Error> {
+    pub fn dkg_end_begin<R: RngCore + CryptoRng>(
+        &mut self,
+        dkg_end_begin: &DkgEndBegin,
+        rng: &mut R,
+    ) -> Result<Vec<Message>, Error> {
         let msgs = vec![];
 
         self.dkg_end_begin_msg = Some(dkg_end_begin.clone());
 
+        for (_, shares) in &dkg_end_begin.dkg_private_shares {
+            let _ = self.dkg_private_shares(shares, rng)?;
+        }
+
         info!(
             signer_id = %self.signer_id,
             dkg_id = %self.dkg_id,
@@ -891,7 +921,9 @@ impl<SignerType: SignerTrait> StateMachine<State, Error> for Signer<SignerType>
 }
 #[cfg(test)]
 /// Test module for signer functionality
-pub mod test {
+mod test {
+    use hashbrown::HashMap;
+
     use crate::{
         common::PolyCommitment,
         curve::{ecdsa, scalar::Scalar},
@@ -926,6 +958,7 @@ pub mod test {
             vec![1],
             Default::default(),
             Default::default(),
+            true,
             &mut rng,
         );
         let public_share = DkgPublicShares {
@@ -963,6 +996,7 @@ pub mod test {
             vec![1],
             Default::default(),
             Default::default(),
+            true,
             &mut rng,
         );
         // publich_shares_done starts out as false
@@ -1001,8 +1035,17 @@ pub mod test {
         public_keys.signers.insert(0, public_key.clone());
         public_keys.key_ids.insert(1, public_key.clone());
 
-        let mut signer =
-            Signer::<SignerType>::new(1, 1, 1, 0, vec![1], private_key, public_keys, &mut rng);
+        let mut signer = Signer::<SignerType>::new(
+            1,
+            1,
+            1,
+            0,
+            vec![1],
+            private_key,
+            public_keys,
+            true,
+            &mut rng,
+        );
         // can_dkg_end starts out as false
         assert!(!signer.can_dkg_end());
 
@@ -1014,10 +1057,17 @@ pub mod test {
         let _ = signer
             .process(&dkg_public_shares[0], &mut rng)
             .expect("failed to process DkgPublicShares");
+        let mut public_shares = HashMap::new();
+        if let Message::DkgPublicShares(shares) = &dkg_public_shares[0] {
+            public_shares.insert(0, shares.clone());
+        } else {
+            panic!("");
+        }
         let dkg_private_begin = Message::DkgPrivateBegin(DkgPrivateBegin {
             dkg_id: 1,
             signer_ids: vec![0],
             key_ids: vec![1],
+            dkg_public_shares: public_shares,
         });
         let dkg_private_shares = signer
             .process(&dkg_private_begin, &mut rng)
@@ -1025,13 +1075,20 @@ pub mod test {
         let _ = signer
             .process(&dkg_private_shares[0], &mut rng)
             .expect("failed to process DkgPrivateShares");
+        let mut private_shares = HashMap::new();
+        if let Message::DkgPrivateShares(shares) = &dkg_private_shares[0] {
+            private_shares.insert(0u32, shares.clone());
+        } else {
+            panic!("");
+        }
         let dkg_end_begin = DkgEndBegin {
             dkg_id: 1,
             signer_ids: vec![0],
             key_ids: vec![1],
+            dkg_private_shares: private_shares,
         };
         let _ = signer
-            .dkg_end_begin(&dkg_end_begin)
+            .dkg_end_begin(&dkg_end_begin, &mut rng)
             .expect("failed to process DkgPrivateShares");
 
         // can_dkg_end should be true
@@ -1054,6 +1111,7 @@ pub mod test {
         .with(fmt::layer())
         .with(EnvFilter::from_default_env())
         .init();*/
+
         let mut rng = create_rng();
         let mut signer = Signer::<SignerType>::new(
             1,
@@ -1063,6 +1121,7 @@ pub mod test {
             vec![1],
             Default::default(),
             Default::default(),
+            true,
             &mut rng,
         );