Add tests for SDK extras and fix execute (#933)

* Don't cancel uploads run without control

Currently `execute` will drop the import handle immediately,
inadvertently canceling the request.

Leak the handle, which will never be used, so that we allow the request
to run to completion.

Closes #930

* Add tests for SDK extras and fix `execute`

We've been relying on the CLI disk import tests to cover the `extras`
disk importer in the SDK. However, the CLI does not use the `execute`
method, and we missed that this was failing immediately due to the
cancellation sender being dropped before the request could awaited.

Add tests for `extras` to validate that both options work. Also, add a
method to DiskImport that does not accept cancellation and use that in
`execute`.

* Add copyright info

Author: wfchandler

Cargo.lock

@@ -27,6 +27,10 @@ toml_edit = { workspace = true }
 uuid = { workspace = true }
 
 [dev-dependencies]
+httpmock = { workspace = true }
+oxide-httpmock = { workspace = true }
+tempfile = { workspace = true }
+test-common = { workspace = true }
 tokio = { workspace = true }
 
 [features]

sdk/Cargo.toml

@@ -99,15 +99,17 @@ pub mod builder {
             self
         }
 
-        /// Return a `Future` for the disk creation and a `DiskImportHandle` which
-        /// can be used to track upload progress and cancel the job.
+        /// Return a `Future` for the disk creation.
         pub fn execute(
             self,
         ) -> Result<
             impl Future<Output = Result<(), DiskImportError>> + 'a,
             Error<crate::types::Error>,
         > {
-            Ok(self.execute_with_control()?.0)
+            let (progress_tx, _progress_rx) = watch::channel(0);
+            let importer = super::types::DiskImport::try_from((self, progress_tx))?;
+
+            Ok(importer.run())
         }
 
         /// Return a `Future` for the disk creation and a `DiskImportHandle` which
@@ -282,6 +284,20 @@ pub mod types {
             Ok(())
         }
 
+        pub async fn run(self) -> Result<(), DiskImportError> {
+            if let Err(e) = self.do_disk_import().await {
+                if let Err(cleanup_err) = self.cleanup().await {
+                    return Err(DiskImportError::Wrapped {
+                        err: cleanup_err.into(),
+                        source: e.into(),
+                    });
+                }
+                return Err(e);
+            }
+
+            Ok(())
+        }
+
         async fn do_disk_import(&self) -> Result<(), DiskImportError> {
             self.check_for_existing_disk().await?;
 

sdk/src/extras/disk.rs

@@ -0,0 +1,141 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+// Copyright 2024 Oxide Computer Company
+
+use httpmock::MockServer;
+use oxide::extras::disk::types::DiskInfo;
+use oxide::extras::ClientExtraDiskExt;
+use oxide::types::Disk;
+use oxide::{Client, ClientConfig};
+use oxide_httpmock::MockServerExt;
+use rand::{thread_rng, Rng, SeedableRng};
+use std::fs;
+use tempfile::TempDir;
+use test_common::JsonMock;
+use uuid::Uuid;
+
+// A disk import where everything succeeds
+#[tokio::test]
+async fn test_disk_import() {
+    let mut src = rand::rngs::SmallRng::seed_from_u64(42);
+    let server = MockServer::start();
+
+    let disk_view_mock = server.disk_view(|when, then| {
+        when.into_inner().any_request();
+        then.client_error(
+            404,
+            &oxide::types::Error {
+                error_code: None,
+                message: "disk not found".into(),
+                request_id: Uuid::mock_value(&mut src).unwrap().to_string(),
+            },
+        );
+    });
+
+    let disk_create_mock = server.disk_create(|when, then| {
+        when.into_inner().any_request();
+        then.created(&Disk {
+            name: "test-import".parse().unwrap(),
+            ..Disk::mock_value(&mut src).unwrap()
+        });
+    });
+
+    let start_bulk_write_mock = server.disk_bulk_write_import_start(|when, then| {
+        when.into_inner().any_request();
+        then.no_content();
+    });
+
+    let disk_bulk_write_mock = server.disk_bulk_write_import(|when, then| {
+        when.into_inner().any_request();
+        then.no_content();
+    });
+
+    let stop_bulk_write_mock = server.disk_bulk_write_import_stop(|when, then| {
+        when.into_inner().any_request();
+        then.no_content();
+    });
+
+    let finalize_mock = server.disk_finalize_import(|when, then| {
+        when.into_inner().any_request();
+        then.no_content();
+    });
+
+    let dir = TempDir::new().unwrap();
+    let image_path = dir.path().join("image.raw");
+
+    let mut rng = thread_rng();
+    let mut contents = vec![0u8; 8192];
+    rng.fill(&mut contents[..]);
+    fs::write(&image_path, contents).unwrap();
+
+    let disk_info = DiskInfo::calculate(image_path.clone(), None, None).unwrap();
+
+    let cfg = ClientConfig::default().with_host_and_token(server.url(""), "test_extra");
+    let client = Client::new_authenticated_config(&cfg).unwrap();
+
+    // Execute without control.
+    client
+        .disk_import()
+        .project("hi")
+        .description("test extra")
+        .upload_thread_ct(1)
+        .disk("test-import")
+        .disk_info(disk_info.clone())
+        .execute()
+        .unwrap()
+        .await
+        .unwrap();
+
+    // Execute with control.
+    let (fut, handle) = client
+        .disk_import()
+        .project("hi")
+        .description("test extra")
+        .upload_thread_ct(8)
+        .disk("test-import")
+        .disk_info(disk_info.clone())
+        .execute_with_control()
+        .unwrap();
+
+    // Confirm progress channel is updated.
+    let mut progress = handle.progress();
+    let progress_updated = tokio::spawn(async move { progress.changed().await.is_ok() });
+    fut.await.unwrap();
+    assert!(progress_updated.await.unwrap());
+
+    // The initial call to `disk_view` races with the cancellation. In general,
+    // the call is sent, but this is likely to vary with available parallelism
+    // and system load. In theory we could see multiple endpoints hit before
+    // the cancel is received, but this seems less likely.
+    //
+    // Assert this endpoint before the next request to avoid defining an
+    // expectation for whether this call should be made upon immediate
+    // cancellation.
+    disk_view_mock.assert_hits(2);
+
+    let (fut_to_cancel, handle) = client
+        .disk_import()
+        .project("hi")
+        .description("test extra")
+        .upload_thread_ct(8)
+        .disk("test-import")
+        .disk_info(disk_info.clone())
+        .execute_with_control()
+        .unwrap();
+
+    // Cancel request immediately.
+    handle.cancel();
+    assert_eq!(
+        fut_to_cancel.await.unwrap_err().to_string(),
+        "Disk import canceled",
+    );
+
+    // No calls received for the cancelled request.
+    disk_create_mock.assert_hits(2);
+    start_bulk_write_mock.assert_hits(2);
+    disk_bulk_write_mock.assert_hits(2);
+    stop_bulk_write_mock.assert_hits(2);
+    finalize_mock.assert_hits(2);
+}