Use TestDatabase outside of nexus-db-queries too

cockroach-admin/src/cockroach_cli.rs

@@ -137,7 +137,7 @@ impl CockroachCli {
 mod tests {
     use super::*;
     use cockroach_admin_types::NodeMembership;
-    use nexus_test_utils::db::test_setup_database;
+    use nexus_test_utils::db::TestDatabase;
     use omicron_test_utils::dev;
     use std::net::SocketAddr;
     use url::Url;
@@ -179,8 +179,8 @@ mod tests {
     #[tokio::test]
     async fn test_node_status_compatibility() {
         let logctx = dev::test_setup_log("test_node_status_compatibility");
-        let mut db = test_setup_database(&logctx.log).await;
-        let db_url = db.listen_url().to_string();
+        let db = TestDatabase::new_without_schema(&logctx.log).await;
+        let db_url = db.crdb().listen_url().to_string();
 
         let expected_headers = "id,address,sql_address,build,started_at,updated_at,locality,is_available,is_live";
 
@@ -225,7 +225,7 @@ mod tests {
         assert_eq!(status[0].is_available, true);
         assert_eq!(status[0].is_live, true);
 
-        db.cleanup().await.unwrap();
+        db.terminate().await;
         logctx.cleanup_successful();
     }
 
@@ -236,8 +236,8 @@ mod tests {
     async fn test_node_decommission_compatibility() {
         let logctx =
             dev::test_setup_log("test_node_decommission_compatibility");
-        let mut db = test_setup_database(&logctx.log).await;
-        let db_url = db.listen_url().to_string();
+        let db = TestDatabase::new_without_schema(&logctx.log).await;
+        let db_url = db.crdb().listen_url().to_string();
 
         let expected_headers =
             "id,is_live,replicas,is_decommissioning,membership,is_draining";
@@ -291,7 +291,7 @@ mod tests {
         assert_eq!(result.is_draining, false);
         assert_eq!(result.notes, &[] as &[&str]);
 
-        db.cleanup().await.unwrap();
+        db.terminate().await;
         logctx.cleanup_successful();
     }
 }

cockroach-admin/src/context.rs

@@ -138,18 +138,19 @@ impl ServerContext {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use nexus_test_utils::db::test_setup_database;
+    use nexus_test_utils::db::TestDatabase;
     use omicron_test_utils::dev;
     use std::net::SocketAddrV6;
     use url::Url;
 
     #[tokio::test]
     async fn test_node_id() {
         let logctx = dev::test_setup_log("test_node_id");
-        let mut db = test_setup_database(&logctx.log).await;
+        let db = TestDatabase::new_without_schema(&logctx.log).await;
+        let crdb = db.crdb();
 
         // Construct a `ServerContext`.
-        let db_url = db.listen_url().to_string();
+        let db_url = crdb.listen_url().to_string();
         let url: Url = db_url.parse().expect("valid url");
         let cockroach_address: SocketAddrV6 = format!(
             "{}:{}",
@@ -173,7 +174,7 @@ mod tests {
 
         // The `OnceCell` should be populated now; even if we shut down the DB,
         // we can still fetch the node ID.
-        db.cleanup().await.unwrap();
+        db.terminate().await;
         let node_id =
             context.node_id().await.expect("successfully read node ID");
         assert_eq!(node_id, "1");

nexus/benches/setup_benchmark.rs

@@ -6,7 +6,7 @@
 
 use criterion::{criterion_group, criterion_main, Criterion};
 use dropshot::test_util::LogContext;
-use nexus_test_utils::db::test_setup_database;
+use nexus_db_queries::db::pub_test_utils::TestDatabase;
 use omicron_test_utils::dev;
 
 // This is the default wrapper around most Nexus integration tests.
@@ -22,8 +22,8 @@ async fn do_full_setup() {
 async fn do_crdb_setup() {
     let cfg = nexus_test_utils::load_test_config();
     let logctx = LogContext::new("crdb_setup", &cfg.pkg.log);
-    let mut db = test_setup_database(&logctx.log).await;
-    db.cleanup().await.unwrap();
+    let db = TestDatabase::new_with_datastore(&logctx.log).await;
+    db.terminate().await;
 }
 
 // Wraps exclusively the ClickhouseDB portion of setup/teardown.

nexus/db-queries/Cargo.toml

@@ -69,7 +69,7 @@ omicron-workspace-hack.workspace = true
 omicron-test-utils = { workspace = true, optional = true }
 
 [features]
-# Enable to export `datastore_test`
+# Enable to export `TestDatabase`
 testing = ["omicron-test-utils"]
 
 [dev-dependencies]

nexus/db-queries/src/db/collection_attach.rs

@@ -577,8 +577,8 @@ where
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::identity::Resource as IdentityResource;
+    use crate::db::pub_test_utils::TestDatabase;
     use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection};
     use chrono::Utc;
     use db_macros::Resource;

nexus/db-queries/src/db/collection_detach.rs

@@ -481,8 +481,8 @@ where
 mod test {
     use super::*;
     use crate::db::collection_attach::DatastoreAttachTarget;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::identity::Resource as IdentityResource;
+    use crate::db::pub_test_utils::TestDatabase;
     use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection};
     use chrono::Utc;
     use db_macros::Resource;

nexus/db-queries/src/db/collection_detach_many.rs

@@ -479,8 +479,8 @@ where
 mod test {
     use super::*;
     use crate::db::collection_attach::DatastoreAttachTarget;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::identity::Resource as IdentityResource;
+    use crate::db::pub_test_utils::TestDatabase;
     use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection};
     use chrono::Utc;
     use db_macros::Resource;

nexus/db-queries/src/db/collection_insert.rs

@@ -405,8 +405,8 @@ where
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::identity::Resource as IdentityResource;
+    use crate::db::pub_test_utils::TestDatabase;
     use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection};
     use chrono::{DateTime, Utc};
     use db_macros::Resource;

nexus/db-queries/src/db/datastore/allow_list.rs

@@ -83,8 +83,8 @@ impl super::DataStore {
 
 #[cfg(test)]
 mod tests {
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::fixed_data::allow_list::USER_FACING_SERVICES_ALLOW_LIST_ID;
+    use crate::db::pub_test_utils::TestDatabase;
     use omicron_common::api::external;
     use omicron_test_utils::dev;
 

nexus/db-queries/src/db/datastore/bgp.rs

@@ -1000,7 +1000,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use omicron_common::api::external::IdentityMetadataCreateParams;
     use omicron_common::api::external::Name;
     use omicron_test_utils::dev;

nexus/db-queries/src/db/datastore/clickhouse_policy.rs

@@ -175,7 +175,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_inventory::now_db_precision;
     use nexus_types::deployment::ClickhouseMode;
     use omicron_test_utils::dev;

nexus/db-queries/src/db/datastore/cockroachdb_node_id.rs

@@ -82,7 +82,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use omicron_test_utils::dev;
 
     #[tokio::test]

nexus/db-queries/src/db/datastore/cockroachdb_settings.rs

@@ -134,7 +134,7 @@ impl DataStore {
 #[cfg(test)]
 mod test {
     use super::CockroachDbSettings;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_types::deployment::CockroachDbClusterVersion;
     use omicron_common::api::external::Error;
     use omicron_test_utils::dev;

nexus/db-queries/src/db/datastore/dataset.rs

@@ -235,7 +235,7 @@ impl DataStore {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_db_model::Generation;
     use nexus_db_model::SledBaseboard;
     use nexus_db_model::SledSystemHardware;

nexus/db-queries/src/db/datastore/db_metadata.rs

@@ -496,7 +496,7 @@ impl DataStore {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use camino::Utf8Path;
     use camino_tempfile::Utf8TempDir;
     use nexus_db_model::SCHEMA_VERSION;

nexus/db-queries/src/db/datastore/deployment.rs

@@ -1730,7 +1730,7 @@ impl RunQueryDsl<DbConnection> for InsertTargetQuery {}
 mod tests {
     use super::*;
 
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use crate::db::raw_query_builder::QueryBuilder;
     use nexus_inventory::now_db_precision;
     use nexus_inventory::CollectionBuilder;

nexus/db-queries/src/db/datastore/deployment/external_networking.rs

@@ -408,7 +408,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use crate::db::queries::ALLOW_FULL_TABLE_SCAN_SQL;
     use anyhow::Context as _;
     use async_bb8_diesel::AsyncSimpleConnection;

nexus/db-queries/src/db/datastore/disk.rs

@@ -843,7 +843,7 @@ impl DataStore {
 mod tests {
     use super::*;
 
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_types::external_api::params;
     use omicron_common::api::external;
     use omicron_test_utils::dev;

nexus/db-queries/src/db/datastore/dns.rs

@@ -729,8 +729,8 @@ impl DataStoreDnsTest for DataStore {
 
 #[cfg(test)]
 mod test {
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::datastore::DnsVersionUpdateBuilder;
+    use crate::db::pub_test_utils::TestDatabase;
     use crate::db::DataStore;
     use crate::db::TransactionError;
     use assert_matches::assert_matches;

nexus/db-queries/src/db/datastore/external_ip.rs

@@ -1132,7 +1132,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_types::deployment::OmicronZoneExternalFloatingIp;
     use nexus_types::deployment::OmicronZoneExternalSnatIp;
     use nexus_types::external_api::shared::IpRange;

nexus/db-queries/src/db/datastore/instance.rs

@@ -1943,10 +1943,10 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::datastore::sled;
     use crate::db::lookup::LookupPath;
     use crate::db::pagination::Paginator;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_db_model::InstanceState;
     use nexus_db_model::Project;
     use nexus_db_model::VmmRuntimeState;

nexus/db-queries/src/db/datastore/inventory.rs

@@ -2444,8 +2444,8 @@ impl DataStoreInventoryTest for DataStore {
 #[cfg(test)]
 mod test {
     use crate::db::datastore::inventory::DataStoreInventoryTest;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::datastore::DataStoreConnection;
+    use crate::db::pub_test_utils::TestDatabase;
     use crate::db::raw_query_builder::{QueryBuilder, TrustedStr};
     use crate::db::schema;
     use crate::db::DataStore;

nexus/db-queries/src/db/datastore/ip_pool.rs

@@ -1121,10 +1121,10 @@ mod test {
     use std::num::NonZeroU32;
 
     use crate::authz;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::model::{
         IpPool, IpPoolResource, IpPoolResourceType, Project,
     };
+    use crate::db::pub_test_utils::TestDatabase;
     use assert_matches::assert_matches;
     use nexus_types::external_api::params;
     use nexus_types::identity::Resource;

nexus/db-queries/src/db/datastore/ipv4_nat_entry.rs

@@ -379,7 +379,7 @@ fn ipv4_nat_next_version() -> diesel::expression::SqlLiteral<BigInt> {
 mod test {
     use std::{net::Ipv4Addr, str::FromStr};
 
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use chrono::Utc;
     use nexus_db_model::{Ipv4NatEntry, Ipv4NatValues, MacAddr, Vni};
     use omicron_common::api::external;

nexus/db-queries/src/db/datastore/migration.rs

@@ -178,9 +178,9 @@ impl DataStore {
 mod tests {
     use super::*;
     use crate::authz;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::lookup::LookupPath;
     use crate::db::model::Instance;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_db_model::Project;
     use nexus_types::external_api::params;
     use nexus_types::silo::DEFAULT_SILO_ID;

nexus/db-queries/src/db/datastore/mod.rs

@@ -78,8 +78,6 @@ mod oximeter;
 mod physical_disk;
 mod probe;
 mod project;
-#[cfg(any(test, feature = "testing"))]
-pub mod pub_test_utils;
 mod quota;
 mod rack;
 mod region;
@@ -446,7 +444,6 @@ mod test {
     use crate::authn;
     use crate::authn::SiloAuthnPolicy;
     use crate::authz;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::datastore::test_utils::{
         IneligibleSledKind, IneligibleSleds,
     };
@@ -459,6 +456,7 @@ mod test {
         Region, SiloUser, SledBaseboard, SledSystemHardware, SledUpdate,
         SshKey, Zpool,
     };
+    use crate::db::pub_test_utils::TestDatabase;
     use crate::db::queries::vpc_subnet::InsertVpcSubnetQuery;
     use chrono::{Duration, Utc};
     use futures::stream;

nexus/db-queries/src/db/datastore/network_interface.rs

@@ -892,7 +892,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_config::NUM_INITIAL_RESERVED_IP_ADDRESSES;
     use nexus_db_fixed_data::vpc_subnet::NEXUS_VPC_SUBNET;
     use omicron_common::address::NEXUS_OPTE_IPV4_SUBNET;

nexus/db-queries/src/db/datastore/oximeter.rs

@@ -292,7 +292,7 @@ impl DataStore {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
+    use crate::db::pub_test_utils::TestDatabase;
     use nexus_types::internal_api::params;
     use omicron_common::api::internal::nexus;
     use omicron_test_utils::dev;

nexus/db-queries/src/db/datastore/physical_disk.rs

@@ -321,12 +321,12 @@ impl DataStore {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::db::datastore::pub_test_utils::TestDatabase;
     use crate::db::datastore::test::{
         sled_baseboard_for_test, sled_system_hardware_for_test,
     };
     use crate::db::lookup::LookupPath;
     use crate::db::model::{PhysicalDiskKind, Sled, SledUpdate};
+    use crate::db::pub_test_utils::TestDatabase;
     use dropshot::PaginationOrder;
     use nexus_db_model::Generation;
     use nexus_sled_agent_shared::inventory::{