From 017d6e897d2bc794ab76c910e4b0d16d1639ba50 Mon Sep 17 00:00:00 2001
From: James MacMahon <james@oxide.computer>
Date: Thu, 19 Sep 2024 14:59:07 -0400
Subject: [PATCH] Add some more resource helpers (#6602)

Add `create_disk_from_snapshot` helper.

Also, when writing unit tests that include expunging a disk, it's still
useful to ask if all Crucible resources were expunged during clean up.
However Nexus will not perform this clean up if a disk is expunged, so
add a few methods to DiskTest to drop a given zpool from the
PerSledDiskState.
---
 nexus/test-utils/src/resource_helpers.rs | 39 ++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/nexus/test-utils/src/resource_helpers.rs b/nexus/test-utils/src/resource_helpers.rs
index 14180459ab..701dac288d 100644
--- a/nexus/test-utils/src/resource_helpers.rs
+++ b/nexus/test-utils/src/resource_helpers.rs
@@ -432,6 +432,28 @@ pub async fn create_disk(
     .await
 }
 
+pub async fn create_disk_from_snapshot(
+    client: &ClientTestContext,
+    project_name: &str,
+    disk_name: &str,
+    snapshot_id: Uuid,
+) -> Disk {
+    let url = format!("/v1/disks?project={}", project_name);
+    object_create(
+        client,
+        &url,
+        &params::DiskCreate {
+            identity: IdentityMetadataCreateParams {
+                name: disk_name.parse().unwrap(),
+                description: String::from("sells rainsticks"),
+            },
+            disk_source: params::DiskSource::Snapshot { snapshot_id },
+            size: ByteCount::from_gibibytes_u32(1),
+        },
+    )
+    .await
+}
+
 pub async fn create_snapshot(
     client: &ClientTestContext,
     project_name: &str,
@@ -1214,6 +1236,10 @@ impl<'a, N: NexusServer> DiskTest<'a, N> {
     }
 
     /// Returns true if all Crucible resources were cleaned up, false otherwise.
+    ///
+    /// Note: be careful performing this test when also peforming physical disk
+    /// expungement, as Nexus will consider resources on those physical disks
+    /// gone and will not attempt to clean them up!
     pub async fn crucible_resources_deleted(&self) -> bool {
         for (sled_id, state) in &self.sleds {
             for zpool in &state.zpools {
@@ -1231,4 +1257,17 @@ impl<'a, N: NexusServer> DiskTest<'a, N> {
 
         true
     }
+
+    /// Drop all of a zpool's resources
+    ///
+    /// Call this before checking `crucible_resources_deleted` if the test has
+    /// also performed a physical disk policy change to "expunged". Nexus will
+    /// _not_ clean up crucible resources on an expunged disk (due to the "gone"
+    /// check that it performs), but it's useful for tests to be able to assert
+    /// all crucible resources are cleaned up.
+    pub async fn remove_zpool(&mut self, zpool_id: Uuid) {
+        for sled in self.sleds.values_mut() {
+            sled.zpools.retain(|zpool| *zpool.id.as_untyped_uuid() != zpool_id);
+        }
+    }
 }