Repository: cloudflare/daphne
Pull request #361 (2023-07-04 20:22:07)

Add benchmarks for `AggregationJobInitReq` handling
This requires a bit of refactoring, so review commit-by-commit. See commit messages for details.


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 0f5db02..6d95f1b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -87,6 +87,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "anes"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
+
 [[package]]
 name = "anstream"
 version = "0.3.2"
@@ -231,6 +237,12 @@ version = "1.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be"
 
+[[package]]
+name = "cast"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
+
 [[package]]
 name = "cc"
 version = "1.0.79"
@@ -304,6 +316,33 @@ dependencies = [
  "phf_codegen",
 ]
 
+[[package]]
+name = "ciborium"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "effd91f6c78e5a4ace8a5d3c0b6bfaec9e2baaef55f3efc00e45fb2e477ee926"
+dependencies = [
+ "ciborium-io",
+ "ciborium-ll",
+ "serde",
+]
+
+[[package]]
+name = "ciborium-io"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdf919175532b369853f5d5e20b26b43112613fd6fe7aee757e35f7a44642656"
+
+[[package]]
+name = "ciborium-ll"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "defaa24ecc093c77630e6c15e17c51f5e187bf35ee514f4e2d67baaa96dae22b"
+dependencies = [
+ "ciborium-io",
+ "half",
+]
+
 [[package]]
 name = "cipher"
 version = "0.3.0"
@@ -423,6 +462,87 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "criterion"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f"
+dependencies = [
+ "anes",
+ "cast",
+ "ciborium",
+ "clap",
+ "criterion-plot",
+ "futures",
+ "is-terminal",
+ "itertools",
+ "num-traits",
+ "once_cell",
+ "oorandom",
+ "plotters",
+ "rayon",
+ "regex",
+ "serde",
+ "serde_derive",
+ "serde_json",
+ "tinytemplate",
+ "tokio",
+ "walkdir",
+]
+
+[[package]]
+name = "criterion-plot"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1"
+dependencies = [
+ "cast",
+ "itertools",
+]
+
+[[package]]
+name = "crossbeam-channel"
+version = "0.5.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a33c2bf77f2df06183c3aa30d1e96c0695a313d4f9c453cc3762a6db39f99200"
+dependencies = [
+ "cfg-if",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef"
+dependencies = [
+ "cfg-if",
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7"
+dependencies = [
+ "autocfg",
+ "cfg-if",
+ "crossbeam-utils",
+ "memoffset",
+ "scopeguard",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "crypto-bigint"
 version = "0.3.2"
@@ -509,6 +629,7 @@ dependencies = [
  "assert_matches",
  "async-trait",
  "base64 0.21.2",
+ "criterion",
  "futures",
  "hex",
  "hpke-rs",
@@ -627,6 +748,12 @@ dependencies = [
  "signature",
 ]
 
+[[package]]
+name = "either"
+version = "1.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91"
+
 [[package]]
 name = "elliptic-curve"
 version = "0.11.12"
@@ -886,6 +1013,12 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "half"
+version = "1.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7"
+
 [[package]]
 name = "hashbrown"
 version = "0.12.3"
@@ -1157,6 +1290,15 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.10.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.6"
@@ -1245,6 +1387,15 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
 
+[[package]]
+name = "memoffset"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"
@@ -1316,6 +1467,12 @@ version = "1.17.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9670a07f94779e00908f3e686eab508878ebb390ba6e604d3a284c00e8d0487b"
 
+[[package]]
+name = "oorandom"
+version = "11.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575"
+
 [[package]]
 name = "opaque-debug"
 version = "0.3.0"
@@ -1526,6 +1683,34 @@ version = "0.3.27"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964"
 
+[[package]]
+name = "plotters"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2c224ba00d7cadd4d5c660deaf2098e5e80e07846537c51f9cfa4be50c1fd45"
+dependencies = [
+ "num-traits",
+ "plotters-backend",
+ "plotters-svg",
+ "wasm-bindgen",
+ "web-sys",
+]
+
+[[package]]
+name = "plotters-backend"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e76628b4d3a7581389a35d5b6e2139607ad7c75b17aed325f210aa91f4a9609"
+
+[[package]]
+name = "plotters-svg"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38f6d39893cca0701371e3c27294f09797214b86f1fb951b89ade8ec04e2abab"
+dependencies = [
+ "plotters-backend",
+]
+
 [[package]]
 name = "poly1305"
 version = "0.7.2"
@@ -1657,6 +1842,28 @@ dependencies = [
  "getrandom",
 ]
 
+[[package]]
+name = "rayon"
+version = "1.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d2df5196e37bcc87abebc0053e20787d73847bb33134a69841207dd0a47f03b"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b8f95bd6966f5c87776639160a66bd8ab9895d9d4ab01ddba9fc60661aebe8d"
+dependencies = [
+ "crossbeam-channel",
+ "crossbeam-deque",
+ "crossbeam-utils",
+ "num_cpus",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.2.16"
@@ -1826,6 +2033,15 @@ version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041"
 
+[[package]]
+name = "same-file"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
+dependencies = [
+ "winapi-util",
+]
+
 [[package]]
 name = "schannel"
 version = "0.1.21"
@@ -2129,6 +2345,16 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "tinytemplate"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "tinyvec"
 version = "1.6.0"
@@ -2405,6 +2631,16 @@ version = "0.9.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
 
+[[package]]
+name = "walkdir"
+version = "2.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36df944cda56c7d8d8b7496af378e6b16de9284591917d307c9b4d313c44e698"
+dependencies = [
+ "same-file",
+ "winapi-util",
+]
+
 [[package]]
 name = "want"
 version = "0.3.0"
@@ -2526,6 +2762,15 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
 
+[[package]]
+name = "winapi-util"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "winapi-x86_64-pc-windows-gnu"
 version = "0.4.0"
diff --git a/Cargo.toml b/Cargo.toml
index 90664e2..6050739 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,6 +16,7 @@ opt-level = "s"
 [workspace.dependencies]
 assert_matches = "1.5.0"
 async-trait = "0.1.68"
+criterion = { version = "0.5.1", features = ["async_tokio"] }
 futures = "0.3.28"
 getrandom = { version = "0.2.9", features = ["js"] } # Required for prio
 hex = { version = "0.4.3", features = ["serde"] }
@@ -30,7 +31,7 @@ ring = "0.16.20"
 serde = { version = "1.0.163", features = ["derive"] }
 serde_json = "1.0.96"
 thiserror = "1.0.40"
-tokio = { version = "1.28.2", features = ["macros", "rt"] }
+tokio = { version = "1.28.2", features = ["macros", "rt-multi-thread"] }
 tracing = "0.1.37"
 url = { version = "2.3.1", features = ["serde"] }
 worker = "0.0.17"
diff --git a/daphne/Cargo.toml b/daphne/Cargo.toml
index 42e7655..6002955 100644
--- a/daphne/Cargo.toml
+++ b/daphne/Cargo.toml
@@ -17,6 +17,10 @@ readme = "../README.md"
 [lib]
 crate-type = ["cdylib", "rlib"]
 
+[[bench]]
+name = "aggregation"
+harness = false
+
 [dependencies]
 assert_matches.workspace = true
 async-trait.workspace = true
@@ -37,6 +41,7 @@ tracing.workspace = true
 url.workspace = true
 
 [dev-dependencies]
+criterion.workspace = true
 matchit.workspace = true
 paste.workspace = true
 tokio.workspace = true
diff --git a/daphne/benches/aggregation.rs b/daphne/benches/aggregation.rs
new file mode 100644
index 0000000..c99d9da
--- /dev/null
+++ b/daphne/benches/aggregation.rs
@@ -0,0 +1,58 @@
+// Copyright (c) 2023 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use daphne::{
+    hpke::HpkeKemId, testing::AggregationJobTest, DapLeaderTransition, DapMeasurement, DapVersion,
+    Prio3Config, VdafConfig,
+};
+
+fn handle_agg_job_init_req(c: &mut Criterion) {
+    let batch_size = 10;
+    let dimension = 100_000;
+
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(4)
+        .enable_all()
+        .build()
+        .unwrap();
+
+    for (vdaf, measurement) in [
+        (
+            VdafConfig::Prio3(Prio3Config::SumVec {
+                bits: 1,
+                len: dimension,
+            }),
+            DapMeasurement::U128Vec(vec![1; dimension]),
+        ),
+        (
+            VdafConfig::Prio2 { dimension },
+            DapMeasurement::U32Vec(vec![1; dimension]),
+        ),
+    ]
+    .into_iter()
+    {
+        let agg_job_test =
+            AggregationJobTest::new(&vdaf, HpkeKemId::X25519HkdfSha256, DapVersion::Draft02);
+
+        let agg_job_init_req = rt.block_on(async {
+            let reports = agg_job_test.produce_reports(vec![measurement; batch_size]);
+            let DapLeaderTransition::Continue(_leader_state, agg_job_init_req) =
+            agg_job_test.produce_agg_job_init_req(reports).await else {
+                panic!("unexpected transition");
+            };
+            agg_job_init_req
+        });
+
+        c.bench_function(&format!("handle_agg_job_init_req {vdaf:?}"), |b| {
+            b.to_async(&rt).iter(|| async {
+                agg_job_test
+                    .handle_agg_job_init_req(&agg_job_init_req)
+                    .await
+            })
+        });
+    }
+}
+
+criterion_group!(benches, handle_agg_job_init_req);
+criterion_main!(benches);
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 6d569e1..60dfd1f 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -53,6 +53,8 @@ use crate::{
     },
 };
 use constants::DapMediaType;
+#[cfg(test)]
+use criterion as _;
 pub use error::DapError;
 use hpke::{HpkeConfig, HpkeKemId};
 use prio::{
@@ -427,7 +429,7 @@ impl AsRef<DapTaskConfig> for DapTaskConfig {
 }
 
 /// A measurement from which a Client generates a report.
-#[derive(Deserialize, Serialize)]
+#[derive(Clone, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum DapMeasurement {
     U64(u64),
@@ -602,8 +604,8 @@ impl DapAggregateShare {
         self.data = None;
     }
 
-    #[cfg(test)]
-    pub(crate) fn try_from_out_shares(
+    /// Aggregate a set of output shares into an aggregate share.
+    pub fn try_from_out_shares(
         out_shares: impl IntoIterator<Item = DapOutputShare>,
     ) -> Result<Self, DapError> {
         let mut agg_share = Self::default();
diff --git a/daphne/src/vdaf/prio3.rs b/daphne/src/vdaf/prio3.rs
index ef03061..5f24a38 100644
--- a/daphne/src/vdaf/prio3.rs
+++ b/daphne/src/vdaf/prio3.rs
@@ -64,7 +64,7 @@ pub(crate) fn prio3_shard(
             let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
             Ok(shard!(vdaf, &measurement, nonce))
         }
-        _ => panic!("prio3_shard: unexpected VDAF config"),
+        _ => panic!("prio3_shard: unexpected VDAF config {config:?}"),
     }
 }
 


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index e198ad4..52c09f2 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -8,17 +8,19 @@ use crate::{
     auth::{BearerToken, BearerTokenProvider},
     constants::DapMediaType,
     fatal_error,
-    hpke::{HpkeConfig, HpkeDecrypter, HpkeReceiverConfig},
+    hpke::{HpkeConfig, HpkeDecrypter, HpkeKemId, HpkeReceiverConfig},
     messages::{
-        AggregationJobId, BatchId, BatchSelector, Collection, CollectionJobId, CollectionReq,
-        Draft02AggregationJobId, HpkeCiphertext, PartialBatchSelector, Report, ReportId,
+        AggregationJobContinueReq, AggregationJobId, AggregationJobInitReq, AggregationJobResp,
+        BatchId, BatchSelector, Collection, CollectionJobId, CollectionReq,
+        Draft02AggregationJobId, HpkeCiphertext, Interval, PartialBatchSelector, Report, ReportId,
         ReportMetadata, TaskId, Time, TransitionFailure,
     },
     metrics::DaphneMetrics,
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    DapAbort, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig,
-    DapHelperState, DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig,
-    DapVersion, MetaAggregationJobId,
+    DapAbort, DapAggregateResult, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError,
+    DapGlobalConfig, DapHelperState, DapHelperTransition, DapLeaderState, DapLeaderTransition,
+    DapLeaderUncommitted, DapMeasurement, DapOutputShare, DapQueryConfig, DapRequest, DapResponse,
+    DapTaskConfig, DapVersion, MetaAggregationJobId, VdafConfig,
 };
 use assert_matches::assert_matches;
 use async_trait::async_trait;
@@ -1061,3 +1063,375 @@ macro_rules! assert_metrics_include {
         }
     }}
 }
+
+/// Scaffolding for testing the aggregation flow.
+pub struct AggregationJobTest {
+    // task parameters
+    pub(crate) task_id: TaskId,
+    pub(crate) task_config: DapTaskConfig,
+    pub(crate) leader_hpke_receiver_config: HpkeReceiverConfig,
+    pub(crate) helper_hpke_receiver_config: HpkeReceiverConfig,
+    pub(crate) client_hpke_config_list: Vec<HpkeConfig>,
+    pub(crate) collector_hpke_receiver_config: HpkeReceiverConfig,
+
+    // aggregation job ID
+    pub(crate) agg_job_id: MetaAggregationJobId<'static>,
+
+    // the current time
+    pub(crate) now: Time,
+
+    // operational parameters
+    #[allow(dead_code)]
+    pub(crate) prometheus_registry: prometheus::Registry,
+    pub(crate) leader_metrics: DaphneMetrics,
+    pub(crate) helper_metrics: DaphneMetrics,
+}
+
+impl AggregationJobTest {
+    /// Create an aggregation job test with the given VDAF config, HPKE KEM algorithm, DAP protocol
+    /// version. The KEM algorithm is used to generate an HPKE config for each party.
+    pub fn new(vdaf: &VdafConfig, kem_id: HpkeKemId, version: DapVersion) -> Self {
+        let mut rng = thread_rng();
+        let now = SystemTime::now()
+            .duration_since(SystemTime::UNIX_EPOCH)
+            .unwrap()
+            .as_secs();
+        let task_id = TaskId(rng.gen());
+        let agg_job_id = MetaAggregationJobId::gen_for_version(&version);
+        let vdaf_verify_key = vdaf.gen_verify_key();
+        let leader_hpke_receiver_config = HpkeReceiverConfig::gen(rng.gen(), kem_id).unwrap();
+        let helper_hpke_receiver_config = HpkeReceiverConfig::gen(rng.gen(), kem_id).unwrap();
+        let collector_hpke_receiver_config = HpkeReceiverConfig::gen(rng.gen(), kem_id).unwrap();
+        let leader_hpke_config = leader_hpke_receiver_config.clone().config;
+        let helper_hpke_config = helper_hpke_receiver_config.clone().config;
+        let collector_hpke_config = collector_hpke_receiver_config.clone().config;
+        let prometheus_registry = prometheus::Registry::new();
+        let leader_metrics =
+            DaphneMetrics::register(&prometheus_registry, Some("test_leader")).unwrap();
+        let helper_metrics =
+            DaphneMetrics::register(&prometheus_registry, Some("test_helper")).unwrap();
+
+        Self {
+            now,
+            task_id,
+            agg_job_id,
+            leader_hpke_receiver_config,
+            helper_hpke_receiver_config,
+            client_hpke_config_list: vec![leader_hpke_config, helper_hpke_config],
+            collector_hpke_receiver_config,
+            task_config: DapTaskConfig {
+                version,
+                leader_url: Url::parse("http://leader.com").unwrap(),
+                helper_url: Url::parse("https://helper.org").unwrap(),
+                time_precision: 500,
+                expiration: now + 500,
+                min_batch_size: 10,
+                query: DapQueryConfig::TimeInterval,
+                vdaf: vdaf.clone(),
+                vdaf_verify_key,
+                collector_hpke_config,
+                taskprov: false,
+            },
+            prometheus_registry,
+            leader_metrics,
+            helper_metrics,
+        }
+    }
+
+    /// For each measurement, generate a report for the given task.
+    ///
+    /// Panics if a measurement is incompatible with the given VDAF.
+    pub fn produce_reports(&self, measurements: Vec<DapMeasurement>) -> Vec<Report> {
+        let mut reports = Vec::with_capacity(measurements.len());
+
+        for measurement in measurements.into_iter() {
+            reports.push(
+                self.task_config
+                    .vdaf
+                    .produce_report(
+                        &self.client_hpke_config_list,
+                        self.now,
+                        &self.task_id,
+                        measurement,
+                        self.task_config.version,
+                    )
+                    .unwrap(),
+            );
+        }
+        reports
+    }
+
+    /// Leader: Produce AggregationJobInitReq.
+    ///
+    /// Panics if the Leader aborts.
+    pub async fn produce_agg_job_init_req(
+        &self,
+        reports: Vec<Report>,
+    ) -> DapLeaderTransition<AggregationJobInitReq> {
+        let metrics = self
+            .leader_metrics
+            .with_host(self.task_config.leader_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .produce_agg_job_init_req(
+                &self.leader_hpke_receiver_config,
+                &self.task_id,
+                &self.task_config,
+                &self.agg_job_id,
+                &PartialBatchSelector::TimeInterval,
+                reports,
+                &metrics,
+            )
+            .await
+            .unwrap()
+    }
+
+    /// Helper: Handle AggregationJobInitReq, produce first AggregationJobResp.
+    ///
+    /// Panics if the Helper aborts.
+    pub async fn handle_agg_job_init_req(
+        &self,
+        agg_job_init_req: &AggregationJobInitReq,
+    ) -> DapHelperTransition<AggregationJobResp> {
+        let metrics = self
+            .helper_metrics
+            .with_host(self.task_config.helper_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .handle_agg_job_init_req(
+                &self.helper_hpke_receiver_config,
+                &self.task_id,
+                &self.task_config,
+                agg_job_init_req,
+                &metrics,
+            )
+            .await
+            .unwrap()
+    }
+
+    /// Leader: Handle first AggregationJobResp, produce AggregationJobContinueReq.
+    ///
+    /// Panics if the Leader aborts.
+    pub fn handle_agg_job_resp(
+        &self,
+        leader_state: DapLeaderState,
+        agg_job_resp: AggregationJobResp,
+    ) -> DapLeaderTransition<AggregationJobContinueReq> {
+        let metrics = self
+            .leader_metrics
+            .with_host(self.task_config.leader_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .handle_agg_job_resp(
+                &self.task_id,
+                &self.agg_job_id,
+                leader_state,
+                agg_job_resp,
+                self.task_config.version,
+                &metrics,
+            )
+            .unwrap()
+    }
+
+    /// Like [`handle_agg_job_resp`] but expect the Leader to abort.
+    pub fn handle_agg_job_resp_expect_err(
+        &self,
+        leader_state: DapLeaderState,
+        agg_job_resp: AggregationJobResp,
+    ) -> DapAbort {
+        let metrics = self
+            .leader_metrics
+            .with_host(self.task_config.leader_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .handle_agg_job_resp(
+                &self.task_id,
+                &self.agg_job_id,
+                leader_state,
+                agg_job_resp,
+                self.task_config.version,
+                &metrics,
+            )
+            .expect_err("handle_agg_job_resp() succeeded; expected failure")
+    }
+
+    /// Helper: Handle AggregationJobContinueReq, produce second AggregationJobResp.
+    ///
+    /// Panics if the Helper aborts.
+    pub fn handle_agg_job_cont_req(
+        &self,
+        helper_state: DapHelperState,
+        agg_job_cont_req: &AggregationJobContinueReq,
+    ) -> DapHelperTransition<AggregationJobResp> {
+        let metrics = self
+            .helper_metrics
+            .with_host(self.task_config.helper_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .handle_agg_job_cont_req(
+                &self.task_id,
+                &self.agg_job_id,
+                helper_state,
+                agg_job_cont_req,
+                &metrics,
+            )
+            .unwrap()
+    }
+
+    /// Like [`handle_agg_job_cont_req`] but expect the Helper to abort.
+    pub fn handle_agg_job_cont_req_expect_err(
+        &self,
+        helper_state: DapHelperState,
+        agg_job_cont_req: &AggregationJobContinueReq,
+    ) -> DapAbort {
+        let metrics = self
+            .helper_metrics
+            .with_host(self.task_config.helper_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .handle_agg_job_cont_req(
+                &self.task_id,
+                &self.agg_job_id,
+                helper_state,
+                agg_job_cont_req,
+                &metrics,
+            )
+            .expect_err("handle_agg_job_cont_req() succeeded; expected failure")
+    }
+
+    /// Leader: Handle the last AggregationJobResp.
+    ///
+    /// Panics if the Leader aborts.
+    pub fn handle_final_agg_job_resp(
+        &self,
+        leader_uncommitted: DapLeaderUncommitted,
+        agg_job_resp: AggregationJobResp,
+    ) -> Vec<DapOutputShare> {
+        let metrics = self
+            .leader_metrics
+            .with_host(self.task_config.leader_url.host_str().unwrap());
+        self.task_config
+            .vdaf
+            .handle_final_agg_job_resp(leader_uncommitted, agg_job_resp, &metrics)
+            .unwrap()
+    }
+
+    /// Produce the Leader's encrypted aggregate share.
+    pub fn produce_leader_encrypted_agg_share(
+        &self,
+        batch_selector: &BatchSelector,
+        agg_share: &DapAggregateShare,
+    ) -> HpkeCiphertext {
+        self.task_config
+            .vdaf
+            .produce_leader_encrypted_agg_share(
+                &self.task_config.collector_hpke_config,
+                &self.task_id,
+                batch_selector,
+                agg_share,
+                self.task_config.version,
+            )
+            .unwrap()
+    }
+
+    /// Produce the Helper's encrypted aggregate share.
+    pub fn produce_helper_encrypted_agg_share(
+        &self,
+        batch_selector: &BatchSelector,
+        agg_share: &DapAggregateShare,
+    ) -> HpkeCiphertext {
+        self.task_config
+            .vdaf
+            .produce_helper_encrypted_agg_share(
+                &self.task_config.collector_hpke_config,
+                &self.task_id,
+                batch_selector,
+                agg_share,
+                self.task_config.version,
+            )
+            .unwrap()
+    }
+
+    /// Collector: Consume the aggregate shares.
+    pub async fn consume_encrypted_agg_shares(
+        &self,
+        batch_selector: &BatchSelector,
+        report_count: u64,
+        enc_agg_shares: Vec<HpkeCiphertext>,
+    ) -> DapAggregateResult {
+        self.task_config
+            .vdaf
+            .consume_encrypted_agg_shares(
+                &self.collector_hpke_receiver_config,
+                &self.task_id,
+                batch_selector,
+                report_count,
+                enc_agg_shares,
+                self.task_config.version,
+            )
+            .await
+            .unwrap()
+    }
+
+    /// Generate a set of reports, aggregate them, and unshard the result.
+    pub async fn roundtrip(&mut self, measurements: Vec<DapMeasurement>) -> DapAggregateResult {
+        let batch_selector = BatchSelector::TimeInterval {
+            batch_interval: Interval {
+                start: self.now,
+                duration: 3600,
+            },
+        };
+
+        // Clients: Shard
+        let reports = self.produce_reports(measurements);
+
+        // Aggregators: Preparation
+        let DapLeaderTransition::Continue(leader_state, agg_job_init_req) = self
+            .produce_agg_job_init_req(reports)
+            .await else {
+                panic!("unexpected transition");
+        };
+        let DapHelperTransition::Continue(helper_state, agg_job_resp) = self
+            .handle_agg_job_init_req(&agg_job_init_req)
+            .await else {
+                panic!("unexpected transition");
+        };
+        let got = DapHelperState::get_decoded(
+            &self.task_config.vdaf,
+            &helper_state
+                .get_encoded(&self.task_config.vdaf)
+                .expect("failed to encode helper state"),
+        )
+        .expect("failed to decode helper state");
+        assert_eq!(got, helper_state);
+
+        let DapLeaderTransition::Uncommitted(uncommitted, agg_cont) = self
+            .handle_agg_job_resp(leader_state, agg_job_resp) else {
+                panic!("unexpected transition");
+            };
+        let DapHelperTransition::Finish(helper_out_shares, agg_job_resp) = self
+            .handle_agg_job_cont_req(helper_state, &agg_cont) else {
+                panic!("unexpected transition");
+            };
+        let leader_out_shares = self.handle_final_agg_job_resp(uncommitted, agg_job_resp);
+        let report_count = u64::try_from(leader_out_shares.len()).unwrap();
+
+        // Leader: Aggregation
+        let leader_agg_share = DapAggregateShare::try_from_out_shares(leader_out_shares).unwrap();
+        let leader_encrypted_agg_share =
+            self.produce_leader_encrypted_agg_share(&batch_selector, &leader_agg_share);
+
+        // Helper: Aggregation
+        let helper_agg_share = DapAggregateShare::try_from_out_shares(helper_out_shares).unwrap();
+        let helper_encrypted_agg_share =
+            self.produce_helper_encrypted_agg_share(&batch_selector, &helper_agg_share);
+
+        // Collector: Unshard
+        self.consume_encrypted_agg_shares(
+            &batch_selector,
+            report_count,
+            vec![leader_encrypted_agg_share, helper_encrypted_agg_share],
+        )
+        .await
+    }
+}
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 3ec6b07..a78beaa 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -5,17 +5,16 @@ use crate::{
     assert_metrics_include, assert_metrics_include_auxiliary_function, async_test_version,
     async_test_versions,
     error::DapAbort,
-    hpke::{HpkeAeadId, HpkeConfig, HpkeKdfId, HpkeKemId, HpkeReceiverConfig},
+    hpke::{HpkeAeadId, HpkeConfig, HpkeKdfId, HpkeKemId},
     messages::{
-        AggregationJobContinueReq, AggregationJobInitReq, AggregationJobResp, BatchSelector,
-        HpkeCiphertext, Interval, PartialBatchSelector, Report, ReportId, ReportShare, TaskId,
-        Time, Transition, TransitionFailure, TransitionVar,
+        AggregationJobInitReq, BatchSelector, Interval, PartialBatchSelector, Report, ReportId,
+        ReportShare, Transition, TransitionFailure, TransitionVar,
     },
-    metrics::DaphneMetrics,
-    test_version, test_versions, DapAggregateResult, DapAggregateShare, DapError, DapHelperState,
-    DapHelperTransition, DapLeaderState, DapLeaderTransition, DapLeaderUncommitted, DapMeasurement,
-    DapOutputShare, DapQueryConfig, DapTaskConfig, DapVersion, MetaAggregationJobId, Prio3Config,
-    VdafAggregateShare, VdafConfig, VdafMessage, VdafState,
+    test_version, test_versions,
+    testing::AggregationJobTest,
+    DapAggregateResult, DapAggregateShare, DapError, DapHelperState, DapHelperTransition,
+    DapLeaderState, DapLeaderTransition, DapLeaderUncommitted, DapMeasurement, DapOutputShare,
+    DapVersion, Prio3Config, VdafAggregateShare, VdafConfig, VdafMessage, VdafState,
 };
 use assert_matches::assert_matches;
 use hpke_rs::HpkePublicKey;
@@ -28,8 +27,7 @@ use prio::{
     },
 };
 use rand::prelude::*;
-use std::{fmt::Debug, time::SystemTime};
-use url::Url;
+use std::fmt::Debug;
 
 impl<M: Debug> DapLeaderTransition<M> {
     pub(crate) fn unwrap_continue(self) -> (DapLeaderState, M) {
@@ -75,7 +73,7 @@ impl<M: Debug> DapHelperTransition<M> {
 const TEST_VDAF: &VdafConfig = &VdafConfig::Prio3(Prio3Config::Count);
 
 async fn roundtrip_report(version: DapVersion) {
-    let t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let report = t
         .task_config
         .vdaf
@@ -153,7 +151,7 @@ async fn roundtrip_report(version: DapVersion) {
 async_test_versions! { roundtrip_report }
 
 fn roundtrip_report_unsupported_hpke_suite(version: DapVersion) {
-    let t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
 
     // The helper's HPKE config indicates a KEM type no supported by the client.
     let unsupported_hpke_config_list = vec![
@@ -183,7 +181,7 @@ fn roundtrip_report_unsupported_hpke_suite(version: DapVersion) {
 test_versions! { roundtrip_report_unsupported_hpke_suite }
 
 async fn produce_agg_job_init_req(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![
         DapMeasurement::U64(1),
         DapMeasurement::U64(0),
@@ -206,7 +204,7 @@ async fn produce_agg_job_init_req(version: DapVersion) {
     }
 
     let (helper_state, agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
     assert_eq!(helper_state.seq.len(), 3);
@@ -219,7 +217,7 @@ async fn produce_agg_job_init_req(version: DapVersion) {
 async_test_versions! { produce_agg_job_init_req }
 
 async fn produce_agg_job_init_req_skip_hpke_decrypt_err(version: DapVersion) {
-    let t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
     // Simulate HPKE decryption error of leader's report share.
@@ -238,7 +236,7 @@ async fn produce_agg_job_init_req_skip_hpke_decrypt_err(version: DapVersion) {
 async_test_versions! { produce_agg_job_init_req_skip_hpke_decrypt_err }
 
 async fn produce_agg_job_init_req_skip_hpke_unknown_config_id(version: DapVersion) {
-    let t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
     // Client tries to send Leader encrypted input with incorrect config ID.
@@ -257,7 +255,7 @@ async fn produce_agg_job_init_req_skip_hpke_unknown_config_id(version: DapVersio
 async_test_versions! { produce_agg_job_init_req_skip_hpke_unknown_config_id }
 
 async fn produce_agg_job_init_req_skip_vdaf_prep_error(version: DapVersion) {
-    let t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = vec![
         t.produce_invalid_report_public_share_decode_failure(DapMeasurement::U64(1), version),
         t.produce_invalid_report_input_share_decode_failure(DapMeasurement::U64(1), version),
@@ -276,17 +274,20 @@ async fn produce_agg_job_init_req_skip_vdaf_prep_error(version: DapVersion) {
 async_test_versions! { produce_agg_job_init_req_skip_vdaf_prep_error }
 
 async fn handle_agg_job_init_req_hpke_decrypt_err(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
     // Simulate HPKE decryption error of helper's report share.
     reports[0].encrypted_input_shares[1].payload[0] ^= 1;
 
-    let (_, agg_req) = t
+    let (_, agg_job_init_req) = t
         .produce_agg_job_init_req(reports.clone())
         .await
         .unwrap_continue();
-    let (_, agg_job_resp) = t.handle_agg_job_init_req(agg_req).await.unwrap_continue();
+    let (_, agg_job_resp) = t
+        .handle_agg_job_init_req(&agg_job_init_req)
+        .await
+        .unwrap_continue();
 
     assert_eq!(agg_job_resp.transitions.len(), 1);
     assert_matches!(
@@ -302,17 +303,20 @@ async fn handle_agg_job_init_req_hpke_decrypt_err(version: DapVersion) {
 async_test_versions! { handle_agg_job_init_req_hpke_decrypt_err }
 
 async fn handle_agg_job_init_req_hpke_unknown_config_id(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
     // Client tries to send Helper encrypted input with incorrect config ID.
     reports[0].encrypted_input_shares[1].config_id ^= 1;
 
-    let (_, agg_req) = t
+    let (_, agg_job_init_req) = t
         .produce_agg_job_init_req(reports.clone())
         .await
         .unwrap_continue();
-    let (_, agg_job_resp) = t.handle_agg_job_init_req(agg_req).await.unwrap_continue();
+    let (_, agg_job_resp) = t
+        .handle_agg_job_init_req(&agg_job_init_req)
+        .await
+        .unwrap_continue();
 
     assert_eq!(agg_job_resp.transitions.len(), 1);
     assert_matches!(
@@ -328,13 +332,13 @@ async fn handle_agg_job_init_req_hpke_unknown_config_id(version: DapVersion) {
 async_test_versions! { handle_agg_job_init_req_hpke_unknown_config_id }
 
 async fn handle_agg_job_init_req_vdaf_prep_error(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let report0 =
         t.produce_invalid_report_public_share_decode_failure(DapMeasurement::U64(1), version);
     let report1 =
         t.produce_invalid_report_input_share_decode_failure(DapMeasurement::U64(1), version);
 
-    let agg_req = AggregationJobInitReq {
+    let agg_job_init_req = AggregationJobInitReq {
         draft02_task_id: t.task_id.for_request_payload(&version),
         draft02_agg_job_id: t.agg_job_id.for_request_payload(),
         agg_param: Vec::new(),
@@ -353,7 +357,10 @@ async fn handle_agg_job_init_req_vdaf_prep_error(version: DapVersion) {
         ],
     };
 
-    let (_, agg_job_resp) = t.handle_agg_job_init_req(agg_req).await.unwrap_continue();
+    let (_, agg_job_resp) = t
+        .handle_agg_job_init_req(&agg_job_init_req)
+        .await
+        .unwrap_continue();
 
     assert_eq!(agg_job_resp.transitions.len(), 2);
     assert_matches!(
@@ -373,12 +380,12 @@ async fn handle_agg_job_init_req_vdaf_prep_error(version: DapVersion) {
 async_test_versions! { handle_agg_job_init_req_vdaf_prep_error }
 
 async fn agg_job_resp_abort_transition_out_of_order(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (_, mut agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -396,12 +403,12 @@ async fn agg_job_resp_abort_transition_out_of_order(version: DapVersion) {
 async_test_versions! { agg_job_resp_abort_transition_out_of_order }
 
 async fn agg_job_resp_abort_report_id_repeated(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (_, mut agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -419,12 +426,12 @@ async_test_versions! { agg_job_resp_abort_report_id_repeated }
 
 async fn agg_job_resp_abort_unrecognized_report_id(version: DapVersion) {
     let mut rng = thread_rng();
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (_, mut agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -443,12 +450,12 @@ async fn agg_job_resp_abort_unrecognized_report_id(version: DapVersion) {
 async_test_versions! { agg_job_resp_abort_unrecognized_report_id }
 
 async fn agg_job_resp_abort_invalid_transition(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (_, mut agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -464,7 +471,7 @@ async fn agg_job_resp_abort_invalid_transition(version: DapVersion) {
 async_test_versions! { agg_job_resp_abort_invalid_transition }
 
 async fn agg_job_cont_req(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![
         DapMeasurement::U64(1),
         DapMeasurement::U64(1),
@@ -475,7 +482,7 @@ async fn agg_job_cont_req(version: DapVersion) {
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (helper_state, agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -528,7 +535,7 @@ async fn agg_job_cont_req(version: DapVersion) {
 async_test_versions! { agg_job_cont_req }
 
 async fn agg_job_cont_req_skip_vdaf_prep_error(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     reports.insert(
         1,
@@ -538,7 +545,7 @@ async fn agg_job_cont_req_skip_vdaf_prep_error(version: DapVersion) {
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (helper_state, agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req.clone())
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -570,12 +577,12 @@ async_test_versions! { agg_job_cont_req_skip_vdaf_prep_error }
 
 async fn agg_cont_abort_unrecognized_report_id(version: DapVersion) {
     let mut rng = thread_rng();
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (helper_state, agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -600,12 +607,12 @@ async fn agg_cont_abort_unrecognized_report_id(version: DapVersion) {
 async_test_versions! { agg_cont_abort_unrecognized_report_id }
 
 async fn agg_job_cont_req_abort_transition_out_of_order(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (helper_state, agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -626,12 +633,12 @@ async fn agg_job_cont_req_abort_transition_out_of_order(version: DapVersion) {
 async_test_versions! { agg_job_cont_req_abort_transition_out_of_order }
 
 async fn agg_job_cont_req_abort_report_id_repeated(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
     let (leader_state, agg_job_init_req) =
         t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (helper_state, agg_job_resp) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -651,7 +658,7 @@ async fn agg_job_cont_req_abort_report_id_repeated(version: DapVersion) {
 async_test_versions! { agg_job_cont_req_abort_report_id_repeated }
 
 async fn encrypted_agg_share(version: DapVersion) {
-    let t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let leader_agg_share = DapAggregateShare {
         report_count: 50,
         min_time: 1637359200,
@@ -695,7 +702,7 @@ async fn encrypted_agg_share(version: DapVersion) {
 async_test_versions! { encrypted_agg_share }
 
 async fn helper_state_serialization(version: DapVersion) {
-    let mut t = Test::new(TEST_VDAF, version);
+    let t = AggregationJobTest::new(TEST_VDAF, HpkeKemId::X25519HkdfSha256, version);
     let reports = t.produce_reports(vec![
         DapMeasurement::U64(1),
         DapMeasurement::U64(1),
@@ -705,7 +712,7 @@ async fn helper_state_serialization(version: DapVersion) {
     ]);
     let (_, agg_job_init_req) = t.produce_agg_job_init_req(reports).await.unwrap_continue();
     let (want, _) = t
-        .handle_agg_job_init_req(agg_job_init_req)
+        .handle_agg_job_init_req(&agg_job_init_req)
         .await
         .unwrap_continue();
 
@@ -718,92 +725,7 @@ async fn helper_state_serialization(version: DapVersion) {
 
 async_test_versions! { helper_state_serialization }
 
-pub(crate) struct Test {
-    now: Time,
-    task_id: TaskId,
-    agg_job_id: MetaAggregationJobId<'static>,
-    task_config: DapTaskConfig,
-    leader_hpke_receiver_config: HpkeReceiverConfig,
-    helper_hpke_receiver_config: HpkeReceiverConfig,
-    client_hpke_config_list: Vec<HpkeConfig>,
-    collector_hpke_receiver_config: HpkeReceiverConfig,
-    prometheus_registry: prometheus::Registry,
-    leader_metrics: DaphneMetrics,
-    helper_metrics: DaphneMetrics,
-}
-
-impl Test {
-    pub(crate) fn new(vdaf: &VdafConfig, version: DapVersion) -> Test {
-        let mut rng = thread_rng();
-        let now = SystemTime::now()
-            .duration_since(SystemTime::UNIX_EPOCH)
-            .unwrap()
-            .as_secs();
-        let task_id = TaskId(rng.gen());
-        let agg_job_id = MetaAggregationJobId::gen_for_version(&version);
-        let vdaf_verify_key = vdaf.gen_verify_key();
-        let leader_hpke_receiver_config =
-            HpkeReceiverConfig::gen(rng.gen(), HpkeKemId::X25519HkdfSha256).unwrap();
-        let helper_hpke_receiver_config =
-            HpkeReceiverConfig::gen(rng.gen(), HpkeKemId::X25519HkdfSha256).unwrap();
-        let collector_hpke_receiver_config =
-            HpkeReceiverConfig::gen(rng.gen(), HpkeKemId::X25519HkdfSha256).unwrap();
-        let leader_hpke_config = leader_hpke_receiver_config.clone().config;
-        let helper_hpke_config = helper_hpke_receiver_config.clone().config;
-        let collector_hpke_config = collector_hpke_receiver_config.clone().config;
-        let prometheus_registry = prometheus::Registry::new();
-        let leader_metrics =
-            DaphneMetrics::register(&prometheus_registry, Some("test_leader")).unwrap();
-        let helper_metrics =
-            DaphneMetrics::register(&prometheus_registry, Some("test_helper")).unwrap();
-
-        Test {
-            now,
-            task_id,
-            agg_job_id,
-            leader_hpke_receiver_config,
-            helper_hpke_receiver_config,
-            client_hpke_config_list: vec![leader_hpke_config, helper_hpke_config],
-            collector_hpke_receiver_config,
-            task_config: DapTaskConfig {
-                version,
-                leader_url: Url::parse("http://leader.com").unwrap(),
-                helper_url: Url::parse("https://helper.org").unwrap(),
-                time_precision: 500,
-                expiration: now + 500,
-                min_batch_size: 10,
-                query: DapQueryConfig::TimeInterval,
-                vdaf: vdaf.clone(),
-                vdaf_verify_key,
-                collector_hpke_config,
-                taskprov: false,
-            },
-            prometheus_registry,
-            leader_metrics,
-            helper_metrics,
-        }
-    }
-
-    fn produce_reports(&self, measurements: Vec<DapMeasurement>) -> Vec<Report> {
-        let mut reports = Vec::with_capacity(measurements.len());
-
-        for measurement in measurements.into_iter() {
-            reports.push(
-                self.task_config
-                    .vdaf
-                    .produce_report(
-                        &self.client_hpke_config_list,
-                        self.now,
-                        &self.task_id,
-                        measurement,
-                        self.task_config.version,
-                    )
-                    .unwrap(),
-            );
-        }
-        reports
-    }
-
+impl AggregationJobTest {
     // Tweak the Helper's share so that decoding succeeds but preparation fails.
     fn produce_invalid_report_vdaf_prep_failure(
         &self,
@@ -888,256 +810,4 @@ impl Test {
             )
             .unwrap()
     }
-
-    async fn produce_agg_job_init_req(
-        &self,
-        reports: Vec<Report>,
-    ) -> DapLeaderTransition<AggregationJobInitReq> {
-        let metrics = self
-            .leader_metrics
-            .with_host(self.task_config.leader_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .produce_agg_job_init_req(
-                &self.leader_hpke_receiver_config,
-                &self.task_id,
-                &self.task_config,
-                &self.agg_job_id,
-                &PartialBatchSelector::TimeInterval,
-                reports,
-                &metrics,
-            )
-            .await
-            .unwrap()
-    }
-
-    async fn handle_agg_job_init_req(
-        &mut self,
-        agg_job_init_req: AggregationJobInitReq,
-    ) -> DapHelperTransition<AggregationJobResp> {
-        let metrics = self
-            .helper_metrics
-            .with_host(self.task_config.helper_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .handle_agg_job_init_req(
-                &self.helper_hpke_receiver_config,
-                &self.task_id,
-                &self.task_config,
-                &agg_job_init_req,
-                &metrics,
-            )
-            .await
-            .unwrap()
-    }
-
-    fn handle_agg_job_resp(
-        &self,
-        leader_state: DapLeaderState,
-        agg_job_resp: AggregationJobResp,
-    ) -> DapLeaderTransition<AggregationJobContinueReq> {
-        let metrics = self
-            .leader_metrics
-            .with_host(self.task_config.leader_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .handle_agg_job_resp(
-                &self.task_id,
-                &self.agg_job_id,
-                leader_state,
-                agg_job_resp,
-                self.task_config.version,
-                &metrics,
-            )
-            .unwrap()
-    }
-
-    fn handle_agg_job_resp_expect_err(
-        &self,
-        leader_state: DapLeaderState,
-        agg_job_resp: AggregationJobResp,
-    ) -> DapAbort {
-        let metrics = self
-            .leader_metrics
-            .with_host(self.task_config.leader_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .handle_agg_job_resp(
-                &self.task_id,
-                &self.agg_job_id,
-                leader_state,
-                agg_job_resp,
-                self.task_config.version,
-                &metrics,
-            )
-            .expect_err("handle_agg_job_resp() succeeded; expected failure")
-    }
-
-    fn handle_agg_job_cont_req(
-        &self,
-        helper_state: DapHelperState,
-        agg_job_cont_req: &AggregationJobContinueReq,
-    ) -> DapHelperTransition<AggregationJobResp> {
-        let metrics = self
-            .helper_metrics
-            .with_host(self.task_config.helper_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .handle_agg_job_cont_req(
-                &self.task_id,
-                &self.agg_job_id,
-                helper_state,
-                agg_job_cont_req,
-                &metrics,
-            )
-            .unwrap()
-    }
-
-    fn handle_agg_job_cont_req_expect_err(
-        &self,
-        helper_state: DapHelperState,
-        agg_job_cont_req: &AggregationJobContinueReq,
-    ) -> DapAbort {
-        let metrics = self
-            .helper_metrics
-            .with_host(self.task_config.helper_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .handle_agg_job_cont_req(
-                &self.task_id,
-                &self.agg_job_id,
-                helper_state,
-                agg_job_cont_req,
-                &metrics,
-            )
-            .expect_err("handle_agg_job_cont_req() succeeded; expected failure")
-    }
-
-    fn handle_final_agg_job_resp(
-        &self,
-        leader_uncommitted: DapLeaderUncommitted,
-        agg_job_resp: AggregationJobResp,
-    ) -> Vec<DapOutputShare> {
-        let metrics = self
-            .leader_metrics
-            .with_host(self.task_config.leader_url.host_str().unwrap());
-        self.task_config
-            .vdaf
-            .handle_final_agg_job_resp(leader_uncommitted, agg_job_resp, &metrics)
-            .unwrap()
-    }
-
-    fn produce_leader_encrypted_agg_share(
-        &self,
-        batch_selector: &BatchSelector,
-        agg_share: &DapAggregateShare,
-    ) -> HpkeCiphertext {
-        self.task_config
-            .vdaf
-            .produce_leader_encrypted_agg_share(
-                &self.task_config.collector_hpke_config,
-                &self.task_id,
-                batch_selector,
-                agg_share,
-                self.task_config.version,
-            )
-            .unwrap()
-    }
-
-    fn produce_helper_encrypted_agg_share(
-        &self,
-        batch_selector: &BatchSelector,
-        agg_share: &DapAggregateShare,
-    ) -> HpkeCiphertext {
-        self.task_config
-            .vdaf
-            .produce_helper_encrypted_agg_share(
-                &self.task_config.collector_hpke_config,
-                &self.task_id,
-                batch_selector,
-                agg_share,
-                self.task_config.version,
-            )
-            .unwrap()
-    }
-
-    async fn consume_encrypted_agg_shares(
-        &self,
-        batch_selector: &BatchSelector,
-        report_count: u64,
-        enc_agg_shares: Vec<HpkeCiphertext>,
-    ) -> DapAggregateResult {
-        self.task_config
-            .vdaf
-            .consume_encrypted_agg_shares(
-                &self.collector_hpke_receiver_config,
-                &self.task_id,
-                batch_selector,
-                report_count,
-                enc_agg_shares,
-                self.task_config.version,
-            )
-            .await
-            .unwrap()
-    }
-
-    pub(crate) async fn roundtrip(
-        &mut self,
-        measurements: Vec<DapMeasurement>,
-    ) -> DapAggregateResult {
-        let batch_selector = BatchSelector::TimeInterval {
-            batch_interval: Interval {
-                start: self.now,
-                duration: 3600,
-            },
-        };
-
-        // Clients: Shard
-        let reports = self.produce_reports(measurements);
-
-        // Aggregators: Preparation
-        let (leader_state, agg_init) = self
-            .produce_agg_job_init_req(reports)
-            .await
-            .unwrap_continue();
-        let (helper_state, agg_job_resp) = self
-            .handle_agg_job_init_req(agg_init)
-            .await
-            .unwrap_continue();
-        let got = DapHelperState::get_decoded(
-            &self.task_config.vdaf,
-            &helper_state
-                .get_encoded(&self.task_config.vdaf)
-                .expect("failed to encode helper state"),
-        )
-        .expect("failed to decode helper state");
-        assert_eq!(got, helper_state);
-
-        let (uncommitted, agg_cont) = self
-            .handle_agg_job_resp(leader_state, agg_job_resp)
-            .unwrap_uncommitted();
-        let (helper_out_shares, agg_job_resp) = self
-            .handle_agg_job_cont_req(helper_state, &agg_cont)
-            .unwrap_finish();
-        let leader_out_shares = self.handle_final_agg_job_resp(uncommitted, agg_job_resp);
-        let report_count = u64::try_from(leader_out_shares.len()).unwrap();
-
-        // Leader: Aggregation
-        let leader_agg_share = DapAggregateShare::try_from_out_shares(leader_out_shares).unwrap();
-        let leader_encrypted_agg_share =
-            self.produce_leader_encrypted_agg_share(&batch_selector, &leader_agg_share);
-
-        // Helper: Aggregation
-        let helper_agg_share = DapAggregateShare::try_from_out_shares(helper_out_shares).unwrap();
-        let helper_encrypted_agg_share =
-            self.produce_helper_encrypted_agg_share(&batch_selector, &helper_agg_share);
-
-        // Collector: Unshard
-        self.consume_encrypted_agg_shares(
-            &batch_selector,
-            report_count,
-            vec![leader_encrypted_agg_share, helper_encrypted_agg_share],
-        )
-        .await
-    }
 }
diff --git a/daphne/src/vdaf/prio2_test.rs b/daphne/src/vdaf/prio2_test.rs
index 7bc5b08..780af16 100644
--- a/daphne/src/vdaf/prio2_test.rs
+++ b/daphne/src/vdaf/prio2_test.rs
@@ -2,13 +2,17 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 use crate::{
-    async_test_version, async_test_versions, vdaf::mod_test::Test, DapAggregateResult,
-    DapMeasurement, DapVersion, VdafConfig,
+    async_test_version, async_test_versions, hpke::HpkeKemId, testing::AggregationJobTest,
+    DapAggregateResult, DapMeasurement, DapVersion, VdafConfig,
 };
 use paste::paste;
 
 async fn roundtrip(version: DapVersion) {
-    let mut t = Test::new(&VdafConfig::Prio2 { dimension: 5 }, version);
+    let mut t = AggregationJobTest::new(
+        &VdafConfig::Prio2 { dimension: 5 },
+        HpkeKemId::X25519HkdfSha256,
+        version,
+    );
     let got = t
         .roundtrip(vec![
             DapMeasurement::U32Vec(vec![1, 1, 0, 0, 1]),