From b51e62e66d760a5144fa4265bd39f9e48508c396 Mon Sep 17 00:00:00 2001
From: Alex Koshelev <koshelev@meta.com>
Date: Tue, 24 Sep 2024 18:18:45 -0700
Subject: [PATCH] Reduce number of iterations for breakdown reveal test

We are having some flakiness with this test: https://github.com/private-attribution/ipa/actions/runs/11018449181/job/30598750956?pr=1307 and I attribute it to having too many iterations (by default it is set to 32). Even running it locally takes a long time and it is not possible to have reliable detection for large routines.
---
 ipa-core/src/helpers/transport/stream/mod.rs         |  2 +-
 ipa-core/src/lib.rs                                  | 12 ++++++------
 .../protocol/ipa_prf/aggregation/breakdown_reveal.rs |  8 ++++++--
 .../ipa_prf/boolean_ops/share_conversion_aby.rs      |  2 +-
 4 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/ipa-core/src/helpers/transport/stream/mod.rs b/ipa-core/src/helpers/transport/stream/mod.rs
index 2f2d6ccc6..59c76cdf4 100644
--- a/ipa-core/src/helpers/transport/stream/mod.rs
+++ b/ipa-core/src/helpers/transport/stream/mod.rs
@@ -187,7 +187,7 @@ mod tests {
             let stream =
                 BodyStream::from_bytes_stream(stream::once(future::ready(Ok(Bytes::from(data)))));
 
-            stream.try_collect::<Vec<_>>().await.unwrap()
+            stream.try_collect::<Vec<_>>().await.unwrap();
         });
     }
 }
diff --git a/ipa-core/src/lib.rs b/ipa-core/src/lib.rs
index a4f625a4c..59cae0106 100644
--- a/ipa-core/src/lib.rs
+++ b/ipa-core/src/lib.rs
@@ -122,10 +122,10 @@ pub(crate) mod test_executor {
 pub(crate) mod test_executor {
     use std::future::Future;
 
-    pub fn run_with<F, Fut, T, const ITER: usize>(f: F) -> T
+    pub fn run_with<F, Fut, const ITER: usize>(f: F)
     where
         F: Fn() -> Fut + Send + Sync + 'static,
-        Fut: Future<Output = T>,
+        Fut: Future<Output = ()>,
     {
         tokio::runtime::Builder::new_multi_thread()
             // enable_all() is common to use to build Tokio runtime, but it enables both IO and time drivers.
@@ -134,16 +134,16 @@ pub(crate) mod test_executor {
             .enable_time()
             .build()
             .unwrap()
-            .block_on(f())
+            .block_on(f());
     }
 
     #[allow(dead_code)]
-    pub fn run<F, Fut, T>(f: F) -> T
+    pub fn run<F, Fut>(f: F)
     where
         F: Fn() -> Fut + Send + Sync + 'static,
-        Fut: Future<Output = T>,
+        Fut: Future<Output = ()>,
     {
-        run_with::<_, _, _, 1>(f)
+        run_with::<_, _, 1>(f);
     }
 }
 
diff --git a/ipa-core/src/protocol/ipa_prf/aggregation/breakdown_reveal.rs b/ipa-core/src/protocol/ipa_prf/aggregation/breakdown_reveal.rs
index cd9ea3eb0..0f5269bbf 100644
--- a/ipa-core/src/protocol/ipa_prf/aggregation/breakdown_reveal.rs
+++ b/ipa-core/src/protocol/ipa_prf/aggregation/breakdown_reveal.rs
@@ -210,7 +210,7 @@ pub mod tests {
         secret_sharing::{
             replicated::semi_honest::AdditiveShare as Replicated, BitDecomposed, TransposeFrom,
         },
-        test_executor::run,
+        test_executor::run_with,
         test_fixture::{Reconstruct, Runner, TestWorld},
     };
 
@@ -224,7 +224,11 @@ pub mod tests {
 
     #[test]
     fn semi_honest_happy_path() {
-        run(|| async {
+        // if shuttle executor is enabled, run this test only once.
+        // it is a very expensive test to explore all possible states,
+        // sometimes github bails after 40 minutes of running it
+        // (workers there are really slow).
+        run_with::<_, _, 3>(|| async {
             let world = TestWorld::default();
             let mut rng = rand::thread_rng();
             let mut expectation = Vec::new();
diff --git a/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs b/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs
index 05e3d3b62..a42bdbbbb 100644
--- a/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs
+++ b/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs
@@ -532,7 +532,7 @@ mod tests {
                     .await
                     .unwrap()
                 })
-                .await
+                .await;
         });
     }