addressed comments. split up the test

mpc/garble/mpc-garble/src/protocol/deap/mod.rs

@@ -871,7 +871,7 @@ impl DEAP {
         Ok(())
     }
 
-    // Returns a reference to the evaluator
+    /// Returns a reference to the evaluator.
     pub(crate) fn ev(&self) -> &Evaluator {
         &self.ev
     }

mpc/garble/mpc-garble/src/protocol/deap/vm.rs

@@ -435,11 +435,11 @@ where
 
 /// This trait provides methods to get peer's encodings.
 trait PeerEncodings {
-    /// Returns the peer's encodings of the provided **input** values.
+    /// Returns the peer's encodings of the provided values.
     ///
     /// # Errors
     ///
-    /// Returns an error if the input value is not found or its encoding is not available.
+    /// Returns an error if the value is not found or its encoding is not available.
     fn get_peer_encodings(
         &self,
         value_ids: &[&str],
@@ -452,12 +452,12 @@ impl<OTS, OTR> PeerEncodings for DEAPVm<OTS, OTR> {
         value_ids: &[&str],
     ) -> Result<Vec<EncodedValue<Active>>, PeerEncodingsError> {
         if self.finalized {
-            return Err(PeerEncodingsError::AlreadyFinalized)?;
+            return Err(PeerEncodingsError::AlreadyFinalized);
         }
 
         let deap = self.deap.as_ref().expect("instance set until finalization");
 
-        Ok(value_ids
+        value_ids
             .iter()
             .map(|id| {
                 // get reference by id
@@ -468,10 +468,10 @@ impl<OTS, OTR> PeerEncodings for DEAPVm<OTS, OTR> {
                 // get encoding by reference
                 match deap.ev().get_encoding(&value_ref) {
                     Some(e) => Ok(e),
-                    None => return Err(PeerEncodingsError::EncodingNotAvailable(value_ref)),
+                    None => Err(PeerEncodingsError::EncodingNotAvailable(value_ref)),
                 }
             })
-            .collect::<Result<Vec<_>, PeerEncodingsError>>()?)
+            .collect::<Result<Vec<_>, PeerEncodingsError>>()
     }
 }
 
@@ -483,8 +483,22 @@ mod tests {
 
     use crate::protocol::deap::mock::create_mock_deap_vm;
 
-    #[tokio::test]
-    async fn test_vm() {
+    use core::future::Future;
+    use mpc_core::Block;
+    use mpc_ot::mock::{MockOTReceiver, MockOTSender};
+
+    // Sets up leader and follower VMs. Returns the VMs and the futures which need to be awaited
+    // to trigger circuit execution.
+    async fn set_up_vms() -> (
+        (
+            DEAPVm<MockOTSender<Block>, MockOTReceiver<Block>>,
+            impl Future<Output = Vec<Value>>,
+        ),
+        (
+            DEAPVm<MockOTSender<Block>, MockOTReceiver<Block>>,
+            impl Future<Output = Vec<Value>>,
+        ),
+    ) {
         let (mut leader_vm, mut follower_vm) = create_mock_deap_vm("test_vm").await;
 
         let mut leader_thread = leader_vm.new_thread("test_thread").await.unwrap();
@@ -502,7 +516,7 @@ mod tests {
                 .unwrap();
             let ciphertext_ref = leader_thread.new_output::<[u8; 16]>("ciphertext").unwrap();
 
-            async {
+            async move {
                 leader_thread
                     .execute(
                         AES128.clone(),
@@ -527,7 +541,7 @@ mod tests {
                 .new_output::<[u8; 16]>("ciphertext")
                 .unwrap();
 
-            async {
+            async move {
                 follower_thread
                     .execute(
                         AES128.clone(),
@@ -541,12 +555,39 @@ mod tests {
             }
         };
 
+        ((leader_vm, leader_fut), (follower_vm, follower_fut))
+    }
+
+    #[tokio::test]
+    async fn test_vm() {
+        let ((mut leader_vm, leader_fut), (mut follower_vm, follower_fut)) = set_up_vms().await;
+
         let (leader_result, follower_result) = futures::join!(leader_fut, follower_fut);
 
         assert_eq!(leader_result, follower_result);
 
-        // These encodings should be available
-        assert!(leader_vm.get_peer_encodings(&["msg", "ciphertext"]).is_ok());
+        let (leader_result, follower_result) =
+            futures::join!(leader_vm.finalize(), follower_vm.finalize());
+
+        leader_result.unwrap();
+        follower_result.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_peer_encodings() {
+        let ((mut leader_vm, leader_fut), (mut follower_vm, follower_fut)) = set_up_vms().await;
+
+        // Encodings are not yet available because the circuit hasn't yet been executed
+        let err = leader_vm.get_peer_encodings(&["msg"]).unwrap_err();
+        assert!(matches!(err, PeerEncodingsError::EncodingNotAvailable(_)));
+
+        // Execute the circuits
+        _ = futures::join!(leader_fut, follower_fut);
+
+        // Encodings must be available now
+        assert!(leader_vm
+            .get_peer_encodings(&["msg", "key", "ciphertext"])
+            .is_ok());
 
         // A non-existent value id will cause an error
         let err = leader_vm
@@ -554,14 +595,9 @@ mod tests {
             .unwrap_err();
         assert!(matches!(err, PeerEncodingsError::ValueIdNotFound(_)));
 
-        let (leader_result, follower_result) =
-            futures::join!(leader_vm.finalize(), follower_vm.finalize());
-
         // Trying to get encodings after finalization will cause an error
+        _ = futures::join!(leader_vm.finalize(), follower_vm.finalize());
         let err = leader_vm.get_peer_encodings(&["msg"]).unwrap_err();
         assert!(matches!(err, PeerEncodingsError::AlreadyFinalized));
-
-        leader_result.unwrap();
-        follower_result.unwrap();
     }
 }