replace test_setup* with a builder (#9446)

Author: davepacheco

nexus/benches/setup_benchmark.rs

@@ -12,9 +12,9 @@ use omicron_test_utils::dev;
 // This is the default wrapper around most Nexus integration tests.
 // Benchmark how long an "empty test" would take.
 async fn do_full_setup() {
-    let ctx =
-        nexus_test_utils::test_setup::<omicron_nexus::Server>("full_setup", 0)
-            .await;
+    let ctx = nexus_test_utils::ControlPlaneBuilder::new("full_setup")
+        .start::<omicron_nexus::Server>()
+        .await;
 
     ctx.teardown().await;
 }

nexus/src/app/sagas/instance_start.rs

@@ -1015,10 +1015,11 @@ mod test {
 
     #[tokio::test]
     async fn should_start_with_dead_switch() {
-        let cptestctx = nexus_test_utils::test_setup::<crate::Server>(
+        let cptestctx = nexus_test_utils::ControlPlaneBuilder::new(
             "should_start_with_dead_switch",
-            3,
         )
+        .with_extra_sled_agents(3)
+        .start::<crate::Server>()
         .await;
 
         let client = &cptestctx.external_client;

nexus/test-utils-macros/src/lib.rs

@@ -126,10 +126,12 @@ pub fn nexus_test(attrs: TokenStream, input: TokenStream) -> TokenStream {
         {
             #input_func
 
-            let ctx = ::nexus_test_utils::test_setup::<#which_nexus>(
+            let ctx = ::nexus_test_utils::ControlPlaneBuilder::new(
                 #func_ident_string,
-                #extra_sled_agents,
-            ).await;
+            )
+            .with_extra_sled_agents(#extra_sled_agents)
+            .start::<#which_nexus>()
+            .await;
             #func_ident(&ctx).await;
             ctx.teardown().await;
         }

nexus/test-utils/src/lib.rs

@@ -25,12 +25,11 @@ pub mod resource_helpers;
 pub mod sql;
 mod starter;
 
+pub use nexus_test::ControlPlaneBuilder;
 pub use nexus_test::ControlPlaneTestContext;
 pub use nexus_test::load_test_config;
 #[cfg(feature = "omicron-dev")]
 pub use nexus_test::omicron_dev_setup_with_config;
-pub use nexus_test::test_setup;
-pub use nexus_test::test_setup_with_config;
 pub use starter::ControlPlaneStarter;
 pub use starter::ControlPlaneTestContextSledAgent;
 pub use starter::register_test_producer;

nexus/test-utils/src/nexus_test.rs

@@ -38,6 +38,64 @@ use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
+pub struct ControlPlaneBuilder<'a> {
+    // required
+    test_name: &'a str,
+
+    // defaults provided by the builder
+    nextra_sled_agents: u16,
+    tls_cert: Option<Certificate>,
+    nexus_config: NexusConfig,
+}
+
+impl<'a> ControlPlaneBuilder<'a> {
+    pub fn new(test_name: &'a str) -> Self {
+        ControlPlaneBuilder {
+            test_name,
+            nextra_sled_agents: 0,
+            tls_cert: None,
+            nexus_config: load_test_config(),
+        }
+    }
+
+    pub fn with_extra_sled_agents(mut self, nextra: u16) -> Self {
+        self.nextra_sled_agents = nextra;
+        self
+    }
+
+    pub fn with_tls_cert(mut self, tls_cert: Option<Certificate>) -> Self {
+        self.tls_cert = tls_cert;
+        self
+    }
+
+    pub fn customize_nexus_config(
+        mut self,
+        f: &dyn Fn(&mut NexusConfig) -> (),
+    ) -> Self {
+        f(&mut self.nexus_config);
+        self
+    }
+
+    pub async fn start<N: NexusServer>(self) -> ControlPlaneTestContext<N> {
+        let mut nexus_config = self.nexus_config;
+        let starter =
+            ControlPlaneStarter::<N>::new(self.test_name, &mut nexus_config);
+        setup_with_config_impl(
+            starter,
+            PopulateCrdb::FromEnvironmentSeed,
+            sim::SimMode::Explicit,
+            self.tls_cert,
+            self.nextra_sled_agents,
+            DEFAULT_SP_SIM_CONFIG.into(),
+            false,
+        )
+        .await
+    }
+}
+
+/// Helper for setting up the control plane for testing and accessing its parts
+///
+/// See [`ControlPlaneBuilder`] for setting one up.
 pub struct ControlPlaneTestContext<N> {
     pub start_time: chrono::DateTime<chrono::Utc>,
     pub external_client: ClientTestContext,
@@ -212,45 +270,7 @@ pub fn load_test_config() -> NexusConfig {
         .expect("failed to load config.test.toml")
 }
 
-pub async fn test_setup<N: NexusServer>(
-    test_name: &str,
-    extra_sled_agents: u16,
-) -> ControlPlaneTestContext<N> {
-    let mut config = load_test_config();
-    test_setup_with_config::<N>(
-        test_name,
-        &mut config,
-        sim::SimMode::Explicit,
-        None,
-        extra_sled_agents,
-        DEFAULT_SP_SIM_CONFIG.into(),
-    )
-    .await
-}
-
-/// Setup routine to use for tests.
-pub async fn test_setup_with_config<N: NexusServer>(
-    test_name: &str,
-    config: &mut NexusConfig,
-    sim_mode: sim::SimMode,
-    initial_cert: Option<Certificate>,
-    extra_sled_agents: u16,
-    gateway_config_file: Utf8PathBuf,
-) -> ControlPlaneTestContext<N> {
-    let starter = ControlPlaneStarter::<N>::new(test_name, config);
-    setup_with_config_impl(
-        starter,
-        PopulateCrdb::FromEnvironmentSeed,
-        sim_mode,
-        initial_cert,
-        extra_sled_agents,
-        gateway_config_file,
-        false,
-    )
-    .await
-}
-
-/// Setup routine to use for `omicron-dev`. Use [`test_setup_with_config`] for
+/// Setup routine to use for `omicron-dev`. Use [`ControlPlaneBuilder`] for
 /// tests.
 ///
 /// The main difference from tests is that this routine ensures the seed tarball

nexus/tests/integration_tests/certificates.rs

@@ -8,13 +8,11 @@ use display_error_chain::ErrorChainExt;
 use dropshot::HttpErrorResponseBody;
 use dropshot::test_util::ClientTestContext;
 use futures::TryStreamExt;
-use gateway_test_utils::setup::DEFAULT_SP_SIM_CONFIG;
 use http::StatusCode;
 use http::method::Method;
 use internal_dns_types::names::DNS_ZONE_EXTERNAL_TESTING;
 use nexus_test_utils::http_testing::AuthnMode;
 use nexus_test_utils::http_testing::NexusRequest;
-use nexus_test_utils::load_test_config;
 use nexus_test_utils::resource_helpers::create_certificate;
 use nexus_test_utils::resource_helpers::delete_certificate;
 use nexus_test_utils_macros::nexus_test;
@@ -345,19 +343,14 @@ async fn test_silo_certificates() {
     let silo3 = SiloCert::new("silo3".parse().unwrap());
 
     // Start Nexus with a TLS server instead of its usual HTTP server.
-    let cptestctx = {
-        let mut config = load_test_config();
-        config.deployment.dropshot_external.tls = true;
-        nexus_test_utils::test_setup_with_config::<omicron_nexus::Server>(
-            "test_silo_certificates",
-            &mut config,
-            omicron_sled_agent::sim::SimMode::Explicit,
-            Some(silo1.cert.clone()),
-            0,
-            DEFAULT_SP_SIM_CONFIG.into(),
-        )
-        .await
-    };
+    let cptestctx =
+        nexus_test_utils::ControlPlaneBuilder::new("test_silo_certificates")
+            .customize_nexus_config(&|config| {
+                config.deployment.dropshot_external.tls = true;
+            })
+            .with_tls_cert(Some(silo1.cert.clone()))
+            .start::<omicron_nexus::Server>()
+            .await;
 
     let nexus_port = cptestctx.external_client.bind_address.port();
 

nexus/tests/integration_tests/console_api.rs

@@ -6,7 +6,6 @@ use anyhow::Context;
 use camino::Utf8PathBuf;
 use dropshot::ResultsPage;
 use dropshot::test_util::ClientTestContext;
-use gateway_test_utils::setup::DEFAULT_SP_SIM_CONFIG;
 use http::{StatusCode, header, method::Method};
 use nexus_auth::context::OpContext;
 use std::env::current_dir;
@@ -22,15 +21,13 @@ use nexus_test_utils::resource_helpers::create_console_session;
 use nexus_test_utils::resource_helpers::{
     create_silo, grant_iam, object_create,
 };
-use nexus_test_utils::{load_test_config, test_setup_with_config};
 use nexus_test_utils_macros::nexus_test;
 use nexus_types::external_api::params::{self, ProjectCreate};
 use nexus_types::external_api::shared::{
     FleetRole, SiloIdentityMode, SiloRole,
 };
 use nexus_types::external_api::{shared, views};
 use omicron_common::api::external::{Error, IdentityMetadataCreateParams};
-use omicron_sled_agent::sim;
 use omicron_test_utils::dev::poll::{CondCheckError, wait_for_condition};
 
 type ControlPlaneTestContext =
@@ -387,20 +384,16 @@ async fn test_assets(cptestctx: &ControlPlaneTestContext) {
 
 #[tokio::test]
 async fn test_absolute_static_dir() {
-    let mut config = load_test_config();
-    config.pkg.console.static_dir =
-        Utf8PathBuf::try_from(current_dir().unwrap())
-            .unwrap()
-            .join("tests/static");
-    let cptestctx = test_setup_with_config::<omicron_nexus::Server>(
-        "test_absolute_static_dir",
-        &mut config,
-        sim::SimMode::Explicit,
-        None,
-        0,
-        DEFAULT_SP_SIM_CONFIG.into(),
-    )
-    .await;
+    let cptestctx =
+        nexus_test_utils::ControlPlaneBuilder::new("test_absolute_static_dir")
+            .customize_nexus_config(&|config| {
+                config.pkg.console.static_dir =
+                    Utf8PathBuf::try_from(current_dir().unwrap())
+                        .unwrap()
+                        .join("tests/static");
+            })
+            .start::<omicron_nexus::Server>()
+            .await;
     let testctx = &cptestctx.external_client;
 
     // existing file is returned
@@ -941,17 +934,14 @@ async fn expect_redirect(testctx: &ClientTestContext, from: &str, to: &str) {
 /// the session was found but is expired. vs not found at all
 #[tokio::test]
 async fn test_session_idle_timeout_deletes_session() {
-    // set idle timeout to 0 so we can test expiration
-    let mut config = load_test_config();
-    config.pkg.console.session_idle_timeout_minutes = 0;
-    let cptestctx = test_setup_with_config::<omicron_nexus::Server>(
+    let cptestctx = nexus_test_utils::ControlPlaneBuilder::new(
         "test_session_idle_timeout_deletes_session",
-        &mut config,
-        sim::SimMode::Explicit,
-        None,
-        0,
-        DEFAULT_SP_SIM_CONFIG.into(),
     )
+    .customize_nexus_config(&|config| {
+        // set idle timeout to 0 so we can test expiration
+        config.pkg.console.session_idle_timeout_minutes = 0;
+    })
+    .start::<omicron_nexus::Server>()
     .await;
     let testctx = &cptestctx.external_client;
 

nexus/tests/integration_tests/device_auth.rs

@@ -7,7 +7,6 @@ use std::num::NonZeroU32;
 use chrono::Utc;
 use dropshot::test_util::ClientTestContext;
 use dropshot::{HttpErrorResponseBody, ResultsPage};
-use gateway_test_utils::setup::DEFAULT_SP_SIM_CONFIG;
 use nexus_auth::authn::USER_TEST_UNPRIVILEGED;
 use nexus_config::NexusConfig;
 use nexus_db_queries::db::fixed_data::silo::DEFAULT_SILO;
@@ -21,7 +20,6 @@ use nexus_test_utils::{
     http_testing::{AuthnMode, NexusRequest, RequestBuilder},
     resource_helpers::grant_iam,
 };
-use nexus_test_utils::{load_test_config, test_setup_with_config};
 use nexus_test_utils_macros::nexus_test;
 use nexus_types::external_api::{params, views};
 use nexus_types::external_api::{
@@ -33,7 +31,6 @@ use nexus_types::external_api::{
 use omicron_uuid_kinds::SiloUserUuid;
 
 use http::{StatusCode, header, method::Method};
-use omicron_sled_agent::sim;
 use oxide_client::types::{FleetRole, SiloRole};
 use serde::Deserialize;
 use tokio::time::{Duration, sleep};
@@ -994,34 +991,38 @@ async fn test_admin_logout_deletes_tokens_and_sessions(
 #[tokio::test]
 async fn test_session_list_excludes_expired() {
     // Test with default TTL - session should not be expired
-    let mut config = load_test_config();
-    test_session_list_with_config(&mut config, 1).await;
+    test_session_list_with_config(&|_config| (), 1).await;
 
     // Test with idle TTL = 0 - session should be expired immediately
-    let mut config = load_test_config();
-    config.pkg.console.session_idle_timeout_minutes = 0;
-    test_session_list_with_config(&mut config, 0).await;
+    test_session_list_with_config(
+        &|config| {
+            config.pkg.console.session_idle_timeout_minutes = 0;
+        },
+        0,
+    )
+    .await;
 
     // Test with abs TTL = 0 - session should be expired immediately
-    let mut config = load_test_config();
-    config.pkg.console.session_absolute_timeout_minutes = 0;
-    test_session_list_with_config(&mut config, 0).await;
+    test_session_list_with_config(
+        &|config| {
+            config.pkg.console.session_absolute_timeout_minutes = 0;
+        },
+        0,
+    )
+    .await;
 }
 
 /// Set up a test context with the given config, create a user in the test suite
 /// silo, and create a session, and assert about the length of the session list
 async fn test_session_list_with_config(
-    config: &mut NexusConfig,
+    modify_config: &dyn Fn(&mut NexusConfig) -> (),
     expected_sessions: usize,
 ) {
-    let cptestctx = test_setup_with_config::<omicron_nexus::Server>(
+    let cptestctx = nexus_test_utils::ControlPlaneBuilder::new(
         "test_session_list_excludes_expired",
-        config,
-        sim::SimMode::Explicit,
-        None,
-        0,
-        DEFAULT_SP_SIM_CONFIG.into(),
     )
+    .customize_nexus_config(modify_config)
+    .start::<omicron_nexus::Server>()
     .await;
     let testctx = &cptestctx.external_client;