tests: Use HttpHandlerTest struct instead of returning tuple (#3286)

* tests: Use HttpHandlerTest struct instead of returning tuple * tests: Hoist up usage of install_test_trace_subscriber()
divviup · Jul 9, 2024 · 5f44b7c · 5f44b7c
1 parent a692ded
commit 5f44b7c
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 70 deletions.
diff --git a/aggregator/src/aggregator.rs b/aggregator/src/aggregator.rs
@@ -3563,6 +3563,7 @@ mod tests {
         where
             R: Runtime + Send + Sync + 'static,
         {
+            install_test_trace_subscriber();
             let clock = MockClock::default();
             let vdaf = Prio3Count::new_count(2).unwrap();
             let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Prio3Count).build();
@@ -3598,8 +3599,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload() {
-        install_test_trace_subscriber();
-
         let UploadTest {
             vdaf,
             aggregator,
@@ -3682,8 +3681,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_batch() {
-        install_test_trace_subscriber();
-
         const BATCH_SIZE: usize = 100;
         let UploadTest {
             vdaf,
@@ -3747,8 +3744,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_wrong_hpke_config_id() {
-        install_test_trace_subscriber();
-
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -3818,8 +3813,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_in_the_future_boundary_condition() {
-        install_test_trace_subscriber();
-
         let UploadTest {
             vdaf,
             aggregator,
@@ -3868,7 +3861,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_in_the_future_past_clock_skew() {
-        install_test_trace_subscriber();
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -3925,8 +3917,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_for_collected_batch() {
-        install_test_trace_subscriber();
-
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -4010,8 +4000,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_encrypted_with_task_specific_key() {
-        install_test_trace_subscriber();
-
         let UploadTest {
             vdaf,
             aggregator,
@@ -4076,7 +4064,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_task_expired() {
-        install_test_trace_subscriber();
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -4137,7 +4124,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_report_expired() {
-        install_test_trace_subscriber();
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -4199,7 +4185,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_faulty_encryption() {
-        install_test_trace_subscriber();
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -4262,7 +4247,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_public_share_decode_failure() {
-        install_test_trace_subscriber();
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,
@@ -4324,7 +4308,6 @@ mod tests {
 
     #[tokio::test]
     async fn upload_report_leader_input_share_decode_failure() {
-        install_test_trace_subscriber();
         let mut runtime_manager = TestRuntimeManager::new();
         let UploadTest {
             aggregator,

diff --git a/aggregator/src/aggregator/aggregation_job_continue.rs b/aggregator/src/aggregator/aggregation_job_continue.rs
@@ -398,7 +398,7 @@ mod tests {
         },
         http_handlers::{
             aggregator_handler,
-            test_util::{setup_http_handler_test, take_problem_details},
+            test_util::{take_problem_details, HttpHandlerTest},
         },
         test_util::default_aggregator_config,
     };
@@ -608,7 +608,12 @@ mod tests {
 
     #[tokio::test]
     async fn leader_rejects_aggregation_job_post() {
-        let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+        let HttpHandlerTest {
+            ephemeral_datastore: _ephemeral_datastore,
+            datastore,
+            handler,
+            ..
+        } = HttpHandlerTest::new().await;
 
         let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Prio3Count).build();
         datastore
@@ -634,7 +639,12 @@ mod tests {
 
     #[tokio::test]
     async fn leader_rejects_aggregation_job_delete() {
-        let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+        let HttpHandlerTest {
+            ephemeral_datastore: _ephemeral_datastore,
+            datastore,
+            handler,
+            ..
+        } = HttpHandlerTest::new().await;
 
         let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Prio3Count).build();
         datastore

diff --git a/aggregator/src/aggregator/http_handlers.rs b/aggregator/src/aggregator/http_handlers.rs
@@ -847,28 +847,4 @@ pub mod test_util {
             }
         }
     }
-
-    pub async fn setup_http_handler_test() -> (
-        MockClock,
-        EphemeralDatastore,
-        Arc<Datastore<MockClock>>,
-        impl Handler,
-    ) {
-        install_test_trace_subscriber();
-        let clock = MockClock::default();
-        let ephemeral_datastore = ephemeral_datastore().await;
-        let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await);
-        datastore.put_global_hpke_key().await.unwrap();
-        let handler = aggregator_handler(
-            datastore.clone(),
-            clock.clone(),
-            TestRuntime::default(),
-            &noop_meter(),
-            default_aggregator_config(),
-        )
-        .await
-        .unwrap();
-
-        (clock, ephemeral_datastore, datastore, handler)
-    }
 }
diff --git a/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs b/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs
@@ -1,8 +1,6 @@
 use crate::aggregator::{
     error::BatchMismatch,
-    http_handlers::test_util::{
-        decode_response_body, setup_http_handler_test, take_problem_details,
-    },
+    http_handlers::test_util::{decode_response_body, take_problem_details, HttpHandlerTest},
 };
 use assert_matches::assert_matches;
 use futures::future::try_join_all;
@@ -52,7 +50,12 @@ pub(crate) async fn post_aggregate_share_request<Q: query_type::QueryType>(
 
 #[tokio::test]
 async fn aggregate_share_request_to_leader() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Fake { rounds: 1 }).build();
@@ -84,7 +87,13 @@ async fn aggregate_share_request_to_leader() {
 
 #[tokio::test]
 async fn aggregate_share_request_invalid_batch_interval() {
-    let (clock, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        clock,
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     const REPORT_EXPIRY_AGE: Duration = Duration::from_seconds(3600);
@@ -143,7 +152,12 @@ async fn aggregate_share_request_invalid_batch_interval() {
 
 #[tokio::test]
 async fn aggregate_share_request() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Fake { rounds: 1 })
         .with_max_batch_query_count(1)

diff --git a/aggregator/src/aggregator/http_handlers/tests/aggregation_job_continue.rs b/aggregator/src/aggregator/http_handlers/tests/aggregation_job_continue.rs
@@ -3,7 +3,7 @@ use crate::aggregator::{
         post_aggregation_job_and_decode, post_aggregation_job_expecting_error,
     },
     empty_batch_aggregations,
-    http_handlers::test_util::{setup_http_handler_test, HttpHandlerTest},
+    http_handlers::test_util::HttpHandlerTest,
     test_util::BATCH_AGGREGATION_SHARD_COUNT,
     tests::generate_helper_report_share,
 };
@@ -1057,7 +1057,12 @@ async fn aggregate_continue_accumulate_batch_aggregation() {
 
 #[tokio::test]
 async fn aggregate_continue_leader_sends_non_continue_or_finish_transition() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Poplar1 { bits: 1 }).build();
@@ -1342,7 +1347,12 @@ async fn aggregate_continue_prep_step_fails() {
 
 #[tokio::test]
 async fn aggregate_continue_unexpected_transition() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Poplar1 { bits: 1 }).build();
@@ -1453,7 +1463,12 @@ async fn aggregate_continue_unexpected_transition() {
 
 #[tokio::test]
 async fn aggregate_continue_out_of_order_transition() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Poplar1 { bits: 1 }).build();
@@ -1623,7 +1638,12 @@ async fn aggregate_continue_out_of_order_transition() {
 
 #[tokio::test]
 async fn aggregate_continue_for_non_waiting_aggregation() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Fake { rounds: 1 }).build();

diff --git a/aggregator/src/aggregator/http_handlers/tests/aggregation_job_init.rs b/aggregator/src/aggregator/http_handlers/tests/aggregation_job_init.rs
@@ -3,9 +3,7 @@ use crate::aggregator::{
     empty_batch_aggregations,
     http_handlers::{
         aggregator_handler,
-        test_util::{
-            decode_response_body, setup_http_handler_test, take_problem_details, HttpHandlerTest,
-        },
+        test_util::{decode_response_body, take_problem_details, HttpHandlerTest},
     },
     test_util::{default_aggregator_config, BATCH_AGGREGATION_SHARD_COUNT},
     tests::{generate_helper_report_share, generate_helper_report_share_for_plaintext},
@@ -49,7 +47,12 @@ use trillium_testing::{assert_headers, prelude::put, TestConn};
 
 #[tokio::test]
 async fn aggregate_leader() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Prio3Count).build();
     datastore
@@ -112,7 +115,12 @@ async fn aggregate_leader() {
 
 #[tokio::test]
 async fn aggregate_wrong_agg_auth_token() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     let dap_auth_token = AuthenticationToken::DapAuth(random());
 

diff --git a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs
@@ -1,8 +1,6 @@
 use crate::aggregator::{
     collection_job_tests::setup_collection_job_test_case,
-    http_handlers::test_util::{
-        decode_response_body, setup_http_handler_test, take_problem_details,
-    },
+    http_handlers::test_util::{decode_response_body, take_problem_details, HttpHandlerTest},
 };
 use janus_aggregator_core::{
     datastore::models::{
@@ -132,7 +130,12 @@ async fn collection_job_put_request_invalid_aggregation_parameter() {
 
 #[tokio::test]
 async fn collection_job_put_request_invalid_batch_size() {
-    let (_, _ephemeral_datastore, datastore, handler) = setup_http_handler_test().await;
+    let HttpHandlerTest {
+        ephemeral_datastore: _ephemeral_datastore,
+        datastore,
+        handler,
+        ..
+    } = HttpHandlerTest::new().await;
 
     // Prepare parameters.
     let task = TaskBuilder::new(QueryType::TimeInterval, VdafInstance::Fake { rounds: 1 })