Make DiskTest easier to use across multiple sleds

dev-tools/reconfigurator-cli/tests/test_basic.rs

@@ -21,7 +21,6 @@ use omicron_test_utils::dev::test_cmds::path_to_executable;
 use omicron_test_utils::dev::test_cmds::redact_variable;
 use omicron_test_utils::dev::test_cmds::run_command;
 use omicron_test_utils::dev::test_cmds::EXIT_SUCCESS;
-use omicron_uuid_kinds::GenericUuid;
 use omicron_uuid_kinds::SledUuid;
 use slog::debug;
 use std::io::BufReader;
@@ -59,19 +58,10 @@ type ControlPlaneTestContext =
 async fn test_blueprint_edit(cptestctx: &ControlPlaneTestContext) {
     // Setup
     //
-    // Add a zpool to both sleds, just to ensure that all new zones can find
+    // Add a zpool to all sleds, just to ensure that all new zones can find
     // a transient filesystem wherever they end up being placed.
-    let _sled_agent_zpools = DiskTestBuilder::new(&cptestctx)
-        .on_sled(SledUuid::from_untyped_uuid(
-            cptestctx.sled_agent.sled_agent.id,
-        ))
-        .with_zpool_count(1)
-        .build()
-        .await;
-    let _sled_agent2_zpools = DiskTestBuilder::new(&cptestctx)
-        .on_sled(SledUuid::from_untyped_uuid(
-            cptestctx.sled_agent2.sled_agent.id,
-        ))
+    DiskTestBuilder::new(&cptestctx)
+        .on_all_sleds()
         .with_zpool_count(1)
         .build()
         .await;

nexus/reconfigurator/execution/src/dns.rs

@@ -503,7 +503,6 @@ mod test {
     use omicron_common::zpool_name::ZpoolName;
     use omicron_test_utils::dev::test_setup_log;
     use omicron_uuid_kinds::ExternalIpUuid;
-    use omicron_uuid_kinds::GenericUuid;
     use omicron_uuid_kinds::OmicronZoneUuid;
     use omicron_uuid_kinds::ZpoolUuid;
     use sled_agent_client::ZoneKind;
@@ -1141,19 +1140,10 @@ mod test {
     async fn test_silos_external_dns_end_to_end(
         cptestctx: &ControlPlaneTestContext,
     ) {
-        // Add a zpool to both sleds, just to ensure that all new zones can find
+        // Add a zpool to all sleds, just to ensure that all new zones can find
         // a transient filesystem wherever they end up being placed.
-        let _sled_agent_zpools = DiskTestBuilder::new(&cptestctx)
-            .on_sled(SledUuid::from_untyped_uuid(
-                cptestctx.sled_agent.sled_agent.id,
-            ))
-            .with_zpool_count(1)
-            .build()
-            .await;
-        let _sled_agent2_zpools = DiskTestBuilder::new(&cptestctx)
-            .on_sled(SledUuid::from_untyped_uuid(
-                cptestctx.sled_agent2.sled_agent.id,
-            ))
+        DiskTestBuilder::new(&cptestctx)
+            .on_all_sleds()
             .with_zpool_count(1)
             .build()
             .await;

nexus/src/app/sagas/disk_create.rs

@@ -829,7 +829,6 @@ pub(crate) mod test {
     use nexus_db_queries::context::OpContext;
     use nexus_db_queries::{authn::saga::Serialized, db::datastore::DataStore};
     use nexus_test_utils::resource_helpers::create_project;
-    use nexus_test_utils::resource_helpers::DiskTest;
     use nexus_test_utils_macros::nexus_test;
     use omicron_common::api::external::ByteCount;
     use omicron_common::api::external::IdentityMetadataCreateParams;
@@ -839,6 +838,8 @@ pub(crate) mod test {
 
     type ControlPlaneTestContext =
         nexus_test_utils::ControlPlaneTestContext<crate::Server>;
+    type DiskTest<'a> =
+        nexus_test_utils::resource_helpers::DiskTest<'a, crate::Server>;
 
     const DISK_NAME: &str = "my-disk";
     const PROJECT_NAME: &str = "springfield-squidport";
@@ -981,9 +982,9 @@ pub(crate) mod test {
 
     async fn no_region_allocations_exist(
         datastore: &DataStore,
-        test: &DiskTest,
+        test: &DiskTest<'_>,
     ) -> bool {
-        for zpool in &test.zpools {
+        for zpool in test.zpools() {
             for dataset in &zpool.datasets {
                 if datastore
                     .regions_total_occupied_size(dataset.id)
@@ -1000,9 +1001,9 @@ pub(crate) mod test {
 
     async fn no_regions_ensured(
         sled_agent: &SledAgent,
-        test: &DiskTest,
+        test: &DiskTest<'_>,
     ) -> bool {
-        for zpool in &test.zpools {
+        for zpool in test.zpools() {
             for dataset in &zpool.datasets {
                 let crucible_dataset =
                     sled_agent.get_crucible_dataset(zpool.id, dataset.id).await;
@@ -1016,7 +1017,7 @@ pub(crate) mod test {
 
     pub(crate) async fn verify_clean_slate(
         cptestctx: &ControlPlaneTestContext,
-        test: &DiskTest,
+        test: &DiskTest<'_>,
     ) {
         let sled_agent = &cptestctx.sled_agent.sled_agent;
         let datastore = cptestctx.server.server_context().nexus.datastore();

nexus/src/app/sagas/region_replacement_start.rs

@@ -799,14 +799,15 @@ pub(crate) mod test {
     use nexus_db_queries::context::OpContext;
     use nexus_test_utils::resource_helpers::create_disk;
     use nexus_test_utils::resource_helpers::create_project;
-    use nexus_test_utils::resource_helpers::DiskTest;
     use nexus_test_utils_macros::nexus_test;
     use nexus_types::identity::Asset;
     use sled_agent_client::types::VolumeConstructionRequest;
     use uuid::Uuid;
 
     type ControlPlaneTestContext =
         nexus_test_utils::ControlPlaneTestContext<crate::Server>;
+    type DiskTest<'a> =
+        nexus_test_utils::resource_helpers::DiskTest<'a, crate::Server>;
 
     const DISK_NAME: &str = "my-disk";
     const PROJECT_NAME: &str = "springfield-squidport";
@@ -816,7 +817,7 @@ pub(crate) mod test {
         cptestctx: &ControlPlaneTestContext,
     ) {
         let mut disk_test = DiskTest::new(cptestctx).await;
-        disk_test.add_zpool_with_dataset(&cptestctx).await;
+        disk_test.add_zpool_with_dataset(cptestctx.first_sled()).await;
 
         let client = &cptestctx.external_client;
         let nexus = &cptestctx.server.server_context().nexus;
@@ -1009,7 +1010,7 @@ pub(crate) mod test {
 
     pub(crate) async fn verify_clean_slate(
         cptestctx: &ControlPlaneTestContext,
-        test: &DiskTest,
+        test: &DiskTest<'_>,
         request: &RegionReplacement,
         affected_volume_original: &Volume,
         affected_region_original: &Region,
@@ -1033,11 +1034,11 @@ pub(crate) mod test {
 
     async fn three_region_allocations_exist(
         datastore: &DataStore,
-        test: &DiskTest,
+        test: &DiskTest<'_>,
     ) -> bool {
         let mut count = 0;
 
-        for zpool in &test.zpools {
+        for zpool in test.zpools() {
             for dataset in &zpool.datasets {
                 if datastore
                     .regions_total_occupied_size(dataset.id)
@@ -1132,7 +1133,7 @@ pub(crate) mod test {
         cptestctx: &ControlPlaneTestContext,
     ) {
         let mut disk_test = DiskTest::new(cptestctx).await;
-        disk_test.add_zpool_with_dataset(&cptestctx).await;
+        disk_test.add_zpool_with_dataset(cptestctx.first_sled()).await;
 
         let log = &cptestctx.logctx.log;
 
@@ -1209,7 +1210,7 @@ pub(crate) mod test {
         cptestctx: &ControlPlaneTestContext,
     ) {
         let mut disk_test = DiskTest::new(cptestctx).await;
-        disk_test.add_zpool_with_dataset(&cptestctx).await;
+        disk_test.add_zpool_with_dataset(cptestctx.first_sled()).await;
 
         let client = &cptestctx.external_client;
         let nexus = &cptestctx.server.server_context().nexus;

nexus/src/app/sagas/snapshot_create.rs

@@ -1617,7 +1617,6 @@ mod test {
     use nexus_test_utils::resource_helpers::create_project;
     use nexus_test_utils::resource_helpers::delete_disk;
     use nexus_test_utils::resource_helpers::object_create;
-    use nexus_test_utils::resource_helpers::DiskTest;
     use nexus_test_utils_macros::nexus_test;
     use nexus_types::external_api::params::InstanceDiskAttachment;
     use omicron_common::api::external::ByteCount;
@@ -1630,6 +1629,9 @@ mod test {
     use sled_agent_client::TestInterfaces as SledAgentTestInterfaces;
     use std::str::FromStr;
 
+    type DiskTest<'a> =
+        nexus_test_utils::resource_helpers::DiskTest<'a, crate::Server>;
+
     #[test]
     fn test_create_snapshot_from_disk_modify_request() {
         let disk = VolumeConstructionRequest::Volume {
@@ -1955,7 +1957,7 @@ mod test {
 
     async fn verify_clean_slate(
         cptestctx: &ControlPlaneTestContext,
-        test: &DiskTest,
+        test: &DiskTest<'_>,
     ) {
         // Verifies:
         // - No disk records exist

nexus/test-utils/src/lib.rs

@@ -130,6 +130,14 @@ pub struct ControlPlaneTestContext<N> {
 }
 
 impl<N: NexusServer> ControlPlaneTestContext<N> {
+    pub fn first_sled(&self) -> SledUuid {
+        SledUuid::from_untyped_uuid(self.sled_agent.sled_agent.id)
+    }
+
+    pub fn all_sled_agents(&self) -> impl Iterator<Item = &sim::Server> {
+        [&self.sled_agent, &self.sled_agent2].into_iter()
+    }
+
     pub fn wildcard_silo_dns_name(&self) -> String {
         format!("*.sys.{}", self.external_dns_zone_name)
     }