diff --git a/tests/testkit/testkit/src/cleanup.rs b/tests/testkit/testkit/src/cleanup.rs index f1a1e0228e..5ae791b7c4 100644 --- a/tests/testkit/testkit/src/cleanup.rs +++ b/tests/testkit/testkit/src/cleanup.rs @@ -18,26 +18,40 @@ pub const CLEANUP_TIMEOUT: Duration = Duration::from_mins(2); type HookFn = Box BoxFuture<'static, Result> + Send>; +/// What a hook is registered for, so it can be registered once and cancelled. +/// The variants are separate namespaces, so keys of different kinds never +/// collide. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum Key { + /// The hook that deletes the sandbox with this name. + Sandbox(String), +} + struct Hook { /// What the hook does, for reporting a failure. name: String, /// Registers the hook at most once; see [`Cleanup::push_keyed`]. - key: Option, + key: Option, run: HookFn, } impl Hook { - fn new(name: String, key: Option, hook: F) -> Self + fn new(name: String, hook: F) -> Self where F: FnOnce() -> Fut + Send + 'static, Fut: Future + Send + 'static, { Self { name, - key, + key: None, run: Box::new(move || hook().boxed()), } } + + fn keyed(mut self, key: Key) -> Self { + self.key = Some(key); + self + } } /// An ordered list of cleanup hooks, shared by a test's context and the @@ -54,25 +68,25 @@ impl Cleanup { F: FnOnce() -> Fut + Send + 'static, Fut: Future + Send + 'static, { - lock(&self.hooks).push(Hook::new(name.into(), None, hook)); + lock(&self.hooks).push(Hook::new(name.into(), hook)); } /// Add a hook unless one with `key` is registered, in which case the /// existing hook keeps its place in the order. - pub(crate) fn push_keyed(&self, key: &str, name: impl Into, hook: F) + pub(crate) fn push_keyed(&self, key: Key, name: impl Into, hook: F) where F: FnOnce() -> Fut + Send + 'static, Fut: Future + Send + 'static, { let mut hooks = lock(&self.hooks); - if !hooks.iter().any(|hook| hook.key.as_deref() == Some(key)) { - hooks.push(Hook::new(name.into(), Some(key.to_string()), hook)); + if !hooks.iter().any(|hook| hook.key.as_ref() == Some(&key)) { + hooks.push(Hook::new(name.into(), hook).keyed(key)); } } /// Drop the hook registered with `key`, if any, without running it. - pub(crate) fn remove(&self, key: &str) { - lock(&self.hooks).retain(|hook| hook.key.as_deref() != Some(key)); + pub(crate) fn remove(&self, key: &Key) { + lock(&self.hooks).retain(|hook| hook.key.as_ref() != Some(key)); } /// Run every hook, last registered first, each limited to `timeout`. A @@ -121,7 +135,11 @@ mod tests { use std::sync::{Arc, Mutex}; use std::time::Duration; - use super::Cleanup; + use super::{Cleanup, Key}; + + fn sandbox(name: &str) -> Key { + Key::Sandbox(name.to_string()) + } use crate::error::Error; const TIMEOUT: Duration = Duration::from_secs(5); @@ -130,7 +148,7 @@ mod tests { /// Register a hook, keyed if `key` is set, that records `name` in `log` /// when it runs. - fn record(cleanup: &Cleanup, key: Option<&str>, log: &Log, name: &'static str) { + fn record(cleanup: &Cleanup, key: Option, log: &Log, name: &'static str) { let log = Arc::clone(log); let hook = async move || { log.lock().unwrap().push(name); @@ -160,12 +178,12 @@ mod tests { async fn keyed_hook_is_registered_once_and_can_be_removed() { let log = Arc::new(Mutex::new(Vec::new())); let cleanup = Cleanup::default(); - record(&cleanup, Some("a"), &log, "a"); + record(&cleanup, Some(sandbox("a")), &log, "a"); record(&cleanup, None, &log, "b"); // Already registered: keeps its place before "b". - record(&cleanup, Some("a"), &log, "a again"); - record(&cleanup, Some("c"), &log, "c"); - cleanup.remove("c"); + record(&cleanup, Some(sandbox("a")), &log, "a again"); + record(&cleanup, Some(sandbox("c")), &log, "c"); + cleanup.remove(&sandbox("c")); assert!(cleanup.run(TIMEOUT).await.is_empty()); assert_eq!(*log.lock().unwrap(), ["b", "a"]); } @@ -174,10 +192,10 @@ mod tests { async fn a_running_hook_can_change_the_remaining_hooks() { let log = Arc::new(Mutex::new(Vec::new())); let cleanup = Cleanup::default(); - record(&cleanup, Some("removed"), &log, "removed"); + record(&cleanup, Some(sandbox("removed")), &log, "removed"); let handle = cleanup.clone(); cleanup.push("remove", async move || { - handle.remove("removed"); + handle.remove(&sandbox("removed")); Ok(()) }); assert!(cleanup.run(TIMEOUT).await.is_empty()); diff --git a/tests/testkit/testkit/src/context.rs b/tests/testkit/testkit/src/context.rs index 3953923b88..c323618224 100644 --- a/tests/testkit/testkit/src/context.rs +++ b/tests/testkit/testkit/src/context.rs @@ -12,7 +12,7 @@ use openshell_sdk::{ }; use rand::Rng; -use crate::cleanup::{CLEANUP_TIMEOUT, Cleanup}; +use crate::cleanup::{CLEANUP_TIMEOUT, Cleanup, Key}; use crate::cli::{Cli, CliOutput}; use crate::error::{Error, Result, ResultExt}; use crate::poll::{Poll, poll_until}; @@ -238,18 +238,22 @@ fn track_sandbox(cleanup: &Cleanup, client: &WorkspaceScopedClient, name: &str) let owned = name.to_string(); // The timeout wraps the whole hook, not just the wait, so a gateway that // never answers the delete cannot hang cleanup. - cleanup.push_keyed(name, format!("delete sandbox {name}"), async move || { - client - .delete_sandbox( - &owned, - DeleteOptions { - allow_missing: true, - }, - ) - .await?; - client.wait_deleted(&owned, CLEANUP_TIMEOUT, None).await?; - Ok(()) - }); + cleanup.push_keyed( + Key::Sandbox(name.to_string()), + format!("delete sandbox {name}"), + async move || { + client + .delete_sandbox( + &owned, + DeleteOptions { + allow_missing: true, + }, + ) + .await?; + client.wait_deleted(&owned, CLEANUP_TIMEOUT, None).await?; + Ok(()) + }, + ); } /// The generated gRPC client, for public API the SDK does not wrap yet. @@ -572,7 +576,7 @@ impl Sandbox { .wait_deleted(&self.name, DELETE_TIMEOUT, Some(&self.id)) .await .context(format!("wait for sandbox {} to be deleted", self.name))?; - self.cleanup.remove(&self.name); + self.cleanup.remove(&Key::Sandbox(self.name.clone())); Ok(()) } }