From 320960f3b1f8cf1d88c2d2484117596dd7435626 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Wed, 16 Oct 2024 17:48:14 -0700 Subject: [PATCH] Use a test wrapper for DB work in db-queries --- nexus/db-queries/src/db/collection_attach.rs | 105 ++++------ nexus/db-queries/src/db/collection_detach.rs | 70 +++---- .../src/db/collection_detach_many.rs | 92 ++++----- nexus/db-queries/src/db/collection_insert.rs | 26 +-- .../db-queries/src/db/datastore/allow_list.rs | 14 +- nexus/db-queries/src/db/datastore/bgp.rs | 10 +- .../src/db/datastore/cockroachdb_node_id.rs | 10 +- .../src/db/datastore/cockroachdb_settings.rs | 18 +- nexus/db-queries/src/db/datastore/dataset.rs | 11 +- .../src/db/datastore/db_metadata.rs | 32 +-- .../db-queries/src/db/datastore/deployment.rs | 45 ++--- .../deployment/external_networking.rs | 17 +- nexus/db-queries/src/db/datastore/disk.rs | 10 +- nexus/db-queries/src/db/datastore/dns.rs | 45 ++--- .../src/db/datastore/external_ip.rs | 10 +- nexus/db-queries/src/db/datastore/instance.rs | 66 +++---- .../db-queries/src/db/datastore/inventory.rs | 31 ++- nexus/db-queries/src/db/datastore/ip_pool.rs | 24 +-- .../src/db/datastore/ipv4_nat_entry.rs | 31 ++- .../db-queries/src/db/datastore/migration.rs | 10 +- nexus/db-queries/src/db/datastore/mod.rs | 139 ++++++------- .../src/db/datastore/network_interface.rs | 10 +- nexus/db-queries/src/db/datastore/oximeter.rs | 51 ++--- .../src/db/datastore/physical_disk.rs | 52 +++-- .../src/db/datastore/pub_test_utils.rs | 116 ++++++++++- nexus/db-queries/src/db/datastore/rack.rs | 52 +++-- .../src/db/datastore/region_replacement.rs | 24 +-- .../datastore/region_snapshot_replacement.rs | 52 +++-- nexus/db-queries/src/db/datastore/saga.rs | 46 ++--- nexus/db-queries/src/db/datastore/sled.rs | 56 +++--- .../src/db/datastore/switch_port.rs | 10 +- .../db-queries/src/db/datastore/test_utils.rs | 6 +- .../virtual_provisioning_collection.rs | 31 ++- nexus/db-queries/src/db/datastore/vmm.rs | 10 +- nexus/db-queries/src/db/datastore/volume.rs | 76 ++++---- .../src/db/datastore/volume_repair.rs | 10 +- nexus/db-queries/src/db/datastore/vpc.rs | 38 ++-- nexus/db-queries/src/db/explain.rs | 18 +- nexus/db-queries/src/db/lookup.rs | 17 +- nexus/db-queries/src/db/pagination.rs | 42 ++-- .../db-queries/src/db/queries/external_ip.rs | 172 +++++++++-------- .../src/db/queries/network_interface.rs | 182 +++++++++--------- nexus/db-queries/src/db/queries/next_item.rs | 60 ++---- nexus/db-queries/src/db/queries/oximeter.rs | 20 +- .../src/db/queries/region_allocation.rs | 11 +- .../virtual_provisioning_collection_update.rs | 38 ++-- nexus/db-queries/src/db/queries/vpc_subnet.rs | 44 +---- nexus/db-queries/src/policy_test/mod.rs | 19 +- nexus/db-queries/src/transaction_retry.rs | 17 +- nexus/metrics-producer-gc/src/lib.rs | 4 +- .../tasks/crdb_node_id_collector.rs | 6 +- 51 files changed, 929 insertions(+), 1177 deletions(-) diff --git a/nexus/db-queries/src/db/collection_attach.rs b/nexus/db-queries/src/db/collection_attach.rs index 066e42fae3..127861b2a8 100644 --- a/nexus/db-queries/src/db/collection_attach.rs +++ b/nexus/db-queries/src/db/collection_attach.rs @@ -577,7 +577,8 @@ where #[cfg(test)] mod test { use super::*; - use crate::db::{self, identity::Resource as IdentityResource}; + use crate::db::datastore::pub_test_utils::TestDatabase; + use crate::db::identity::Resource as IdentityResource; use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection}; use chrono::Utc; use db_macros::Resource; @@ -585,7 +586,6 @@ mod test { use diesel::pg::Pg; use diesel::QueryDsl; use diesel::SelectableHelper; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external::{IdentityMetadataCreateParams, Name}; use omicron_test_utils::dev; use uuid::Uuid; @@ -869,11 +869,9 @@ mod test { async fn test_attach_missing_collection_fails() { let logctx = dev::test_setup_log("test_attach_missing_collection_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -891,17 +889,15 @@ mod test { assert!(matches!(attach, Err(AttachError::CollectionNotFound))); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_missing_resource_fails() { let logctx = dev::test_setup_log("test_attach_missing_resource_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = setup_db(&pool).await; @@ -929,17 +925,15 @@ mod test { // The collection should remain unchanged. assert_eq!(collection, get_collection(collection_id, &conn).await); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_once() { let logctx = dev::test_setup_log("test_attach_once"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = setup_db(&pool).await; @@ -978,17 +972,15 @@ mod test { ); assert_eq!(returned_resource, get_resource(resource_id, &conn).await); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_once_synchronous() { let logctx = dev::test_setup_log("test_attach_once_synchronous"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = setup_db(&pool).await; @@ -1028,19 +1020,16 @@ mod test { ); assert_eq!(returned_resource, get_resource(resource_id, &conn).await); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_multiple_times() { let logctx = dev::test_setup_log("test_attach_multiple_times"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; const RESOURCE_COUNT: u32 = 5; @@ -1085,19 +1074,16 @@ mod test { ); } - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_beyond_capacity_fails() { let logctx = dev::test_setup_log("test_attach_beyond_capacity_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -1150,19 +1136,16 @@ mod test { _ => panic!("Unexpected error: {:?}", err), }; - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_while_already_attached() { let logctx = dev::test_setup_log("test_attach_while_already_attached"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -1258,19 +1241,16 @@ mod test { _ => panic!("Unexpected error: {:?}", err), }; - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_with_filters() { let logctx = dev::test_setup_log("test_attach_once"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -1314,19 +1294,16 @@ mod test { assert_eq!(returned_resource, get_resource(resource_id, &conn).await); assert_eq!(returned_resource.description(), "new description"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_deleted_resource_fails() { let logctx = dev::test_setup_log("test_attach_deleted_resource_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -1360,19 +1337,16 @@ mod test { .await; assert!(matches!(attach, Err(AttachError::ResourceNotFound))); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_attach_without_update_filter() { let logctx = dev::test_setup_log("test_attach_without_update_filter"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -1417,8 +1391,7 @@ mod test { .collection_id .is_none()); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/collection_detach.rs b/nexus/db-queries/src/db/collection_detach.rs index ce41545487..cdf8e111c7 100644 --- a/nexus/db-queries/src/db/collection_detach.rs +++ b/nexus/db-queries/src/db/collection_detach.rs @@ -481,7 +481,8 @@ where mod test { use super::*; use crate::db::collection_attach::DatastoreAttachTarget; - use crate::db::{self, identity::Resource as IdentityResource}; + use crate::db::datastore::pub_test_utils::TestDatabase; + use crate::db::identity::Resource as IdentityResource; use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection}; use chrono::Utc; use db_macros::Resource; @@ -489,7 +490,6 @@ mod test { use diesel::pg::Pg; use diesel::QueryDsl; use diesel::SelectableHelper; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external::{IdentityMetadataCreateParams, Name}; use omicron_test_utils::dev; use uuid::Uuid; @@ -782,11 +782,9 @@ mod test { async fn test_detach_missing_collection_fails() { let logctx = dev::test_setup_log("test_detach_missing_collection_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -803,19 +801,16 @@ mod test { assert!(matches!(detach, Err(DetachError::CollectionNotFound))); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_missing_resource_fails() { let logctx = dev::test_setup_log("test_detach_missing_resource_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -840,19 +835,16 @@ mod test { // The collection should remain unchanged. assert_eq!(collection, get_collection(collection_id, &conn).await); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_once() { let logctx = dev::test_setup_log("test_detach_once"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -881,19 +873,16 @@ mod test { // The returned value should be the latest value in the DB. assert_eq!(returned_resource, get_resource(resource_id, &conn).await); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_while_already_detached() { let logctx = dev::test_setup_log("test_detach_while_already_detached"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -946,19 +935,16 @@ mod test { _ => panic!("Unexpected error: {:?}", err), }; - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_deleted_resource_fails() { let logctx = dev::test_setup_log("test_detach_deleted_resource_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -991,19 +977,16 @@ mod test { .await; assert!(matches!(detach, Err(DetachError::ResourceNotFound))); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_without_update_filter() { let logctx = dev::test_setup_log("test_detach_without_update_filter"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -1049,8 +1032,7 @@ mod test { .collection_id .is_some()); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/collection_detach_many.rs b/nexus/db-queries/src/db/collection_detach_many.rs index eb1a67ffff..15267dd7ee 100644 --- a/nexus/db-queries/src/db/collection_detach_many.rs +++ b/nexus/db-queries/src/db/collection_detach_many.rs @@ -479,7 +479,8 @@ where mod test { use super::*; use crate::db::collection_attach::DatastoreAttachTarget; - use crate::db::{self, identity::Resource as IdentityResource}; + use crate::db::datastore::pub_test_utils::TestDatabase; + use crate::db::identity::Resource as IdentityResource; use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection}; use chrono::Utc; use db_macros::Resource; @@ -487,7 +488,6 @@ mod test { use diesel::pg::Pg; use diesel::QueryDsl; use diesel::SelectableHelper; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external::{IdentityMetadataCreateParams, Name}; use omicron_test_utils::dev; use uuid::Uuid; @@ -774,11 +774,9 @@ mod test { async fn test_detach_missing_collection_fails() { let logctx = dev::test_setup_log("test_detach_missing_collection_fails"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let _resource_id = uuid::Uuid::new_v4(); @@ -796,8 +794,7 @@ mod test { assert!(matches!(detach, Err(DetachManyError::CollectionNotFound))); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -805,11 +802,9 @@ mod test { async fn test_detach_missing_resource_succeeds() { let logctx = dev::test_setup_log("test_detach_missing_resource_succeeds"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let _resource_id = uuid::Uuid::new_v4(); @@ -839,19 +834,16 @@ mod test { get_collection(collection_id, &conn).await ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_once() { let logctx = dev::test_setup_log("test_detach_once"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -883,19 +875,16 @@ mod test { get_collection(collection_id, &conn).await ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_once_synchronous() { let logctx = dev::test_setup_log("test_detach_once_synchronous"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -929,19 +918,16 @@ mod test { get_collection(collection_id, &conn).await ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_while_already_detached() { let logctx = dev::test_setup_log("test_detach_while_already_detached"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -986,19 +972,16 @@ mod test { "... and again!" ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_filter_collection() { let logctx = dev::test_setup_log("test_detach_filter_collection"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); @@ -1038,19 +1021,16 @@ mod test { _ => panic!("Unexpected error: {:?}", err), }; - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_deleted_resource() { let logctx = dev::test_setup_log("test_detach_deleted_resource"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -1097,19 +1077,16 @@ mod test { &collection_id, ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_detach_many() { let logctx = dev::test_setup_log("test_detach_many"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; // Create the collection and some resources. let collection_id1 = uuid::Uuid::new_v4(); @@ -1173,8 +1150,7 @@ mod test { &collection_id2 ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/collection_insert.rs b/nexus/db-queries/src/db/collection_insert.rs index fa1c17f789..7f8275e594 100644 --- a/nexus/db-queries/src/db/collection_insert.rs +++ b/nexus/db-queries/src/db/collection_insert.rs @@ -405,14 +405,14 @@ where #[cfg(test)] mod test { use super::*; - use crate::db::{self, identity::Resource as IdentityResource}; + use crate::db::datastore::pub_test_utils::TestDatabase; + use crate::db::identity::Resource as IdentityResource; use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection}; use chrono::{DateTime, Utc}; use db_macros::Resource; use diesel::expression_methods::ExpressionMethods; use diesel::pg::Pg; use diesel::QueryDsl; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; table! { @@ -556,11 +556,9 @@ mod test { #[tokio::test] async fn test_collection_not_present() { let logctx = dev::test_setup_log("test_collection_not_present"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -579,19 +577,16 @@ mod test { .await; assert!(matches!(insert, Err(AsyncInsertError::CollectionNotFound))); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_collection_present() { let logctx = dev::test_setup_log("test_collection_present"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - - let conn = setup_db(&pool).await; + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = setup_db(pool).await; let collection_id = uuid::Uuid::new_v4(); let resource_id = uuid::Uuid::new_v4(); @@ -643,8 +638,7 @@ mod test { // Make sure rcgen got incremented assert_eq!(collection_rcgen, 2); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/allow_list.rs b/nexus/db-queries/src/db/datastore/allow_list.rs index dbe13aafaa..ce839ebcbc 100644 --- a/nexus/db-queries/src/db/datastore/allow_list.rs +++ b/nexus/db-queries/src/db/datastore/allow_list.rs @@ -83,19 +83,16 @@ impl super::DataStore { #[cfg(test)] mod tests { - use crate::db::{ - datastore::test_utils::datastore_test, - fixed_data::allow_list::USER_FACING_SERVICES_ALLOW_LIST_ID, - }; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; + use crate::db::fixed_data::allow_list::USER_FACING_SERVICES_ALLOW_LIST_ID; use omicron_common::api::external; use omicron_test_utils::dev; #[tokio::test] async fn test_allowed_source_ip_database_ops() { let logctx = dev::test_setup_log("test_allowed_source_ip_database_ops"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Should have the default to start with. let record = datastore @@ -203,8 +200,7 @@ mod tests { "Updated allowed IPs are incorrect" ); - datastore.terminate().await; - db.cleanup().await.expect("failed to cleanup database"); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/bgp.rs b/nexus/db-queries/src/db/datastore/bgp.rs index eb7814ff9b..ccf5c6bb75 100644 --- a/nexus/db-queries/src/db/datastore/bgp.rs +++ b/nexus/db-queries/src/db/datastore/bgp.rs @@ -1000,8 +1000,7 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use omicron_common::api::external::IdentityMetadataCreateParams; use omicron_common::api::external::Name; use omicron_test_utils::dev; @@ -1011,8 +1010,8 @@ mod tests { let logctx = dev::test_setup_log( "test_delete_bgp_config_and_announce_set_by_name", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let config_name: Name = "testconfig47".parse().unwrap(); let announce_name: Name = "testannounce47".parse().unwrap(); @@ -1069,8 +1068,7 @@ mod tests { .await .expect("delete announce set by name"); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/cockroachdb_node_id.rs b/nexus/db-queries/src/db/datastore/cockroachdb_node_id.rs index 233194e0f4..1c1a699c26 100644 --- a/nexus/db-queries/src/db/datastore/cockroachdb_node_id.rs +++ b/nexus/db-queries/src/db/datastore/cockroachdb_node_id.rs @@ -82,16 +82,15 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use omicron_test_utils::dev; #[tokio::test] async fn test_cockroachdb_node_id() { let logctx = dev::test_setup_log("test_service_network_interfaces_list"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Make up a CRDB zone id. let crdb_zone_id = OmicronZoneUuid::new_v4(); @@ -160,8 +159,7 @@ mod tests { .expect("looked up node ID"); assert_eq!(node_id.as_deref(), Some(fake_node_id)); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/cockroachdb_settings.rs b/nexus/db-queries/src/db/datastore/cockroachdb_settings.rs index 5994abdd2e..ba7c302f83 100644 --- a/nexus/db-queries/src/db/datastore/cockroachdb_settings.rs +++ b/nexus/db-queries/src/db/datastore/cockroachdb_settings.rs @@ -133,24 +133,17 @@ impl DataStore { #[cfg(test)] mod test { - use super::{CockroachDbSettings, OpContext}; - use nexus_test_utils::db::test_setup_database; + use super::CockroachDbSettings; + use crate::db::datastore::pub_test_utils::TestDatabase; use nexus_types::deployment::CockroachDbClusterVersion; use omicron_common::api::external::Error; use omicron_test_utils::dev; - use std::sync::Arc; #[tokio::test] async fn test_preserve_downgrade() { let logctx = dev::test_setup_log("test_preserve_downgrade"); - let mut db = test_setup_database(&logctx.log).await; - let (_, datastore) = - crate::db::datastore::test_utils::datastore_test(&logctx, &db) - .await; - let opctx = OpContext::for_tests( - logctx.log.new(o!()), - Arc::clone(&datastore) as Arc, - ); + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let settings = datastore.cockroachdb_settings(&opctx).await.unwrap(); let version: CockroachDbClusterVersion = @@ -247,8 +240,7 @@ mod test { } } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/dataset.rs b/nexus/db-queries/src/db/datastore/dataset.rs index a764f8ca70..1843df0c7d 100644 --- a/nexus/db-queries/src/db/datastore/dataset.rs +++ b/nexus/db-queries/src/db/datastore/dataset.rs @@ -235,21 +235,19 @@ impl DataStore { #[cfg(test)] mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use nexus_db_model::Generation; use nexus_db_model::SledBaseboard; use nexus_db_model::SledSystemHardware; use nexus_db_model::SledUpdate; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::internal::shared::DatasetKind as ApiDatasetKind; use omicron_test_utils::dev; #[tokio::test] async fn test_insert_if_not_exists() { let logctx = dev::test_setup_log("inventory_insert"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; - let opctx = &opctx; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // There should be no datasets initially. assert_eq!( @@ -383,8 +381,7 @@ mod test { expected_datasets, ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/db_metadata.rs b/nexus/db-queries/src/db/datastore/db_metadata.rs index 1277de1983..fbb6cd35c8 100644 --- a/nexus/db-queries/src/db/datastore/db_metadata.rs +++ b/nexus/db-queries/src/db/datastore/db_metadata.rs @@ -496,32 +496,26 @@ impl DataStore { #[cfg(test)] mod test { use super::*; + use crate::db::datastore::pub_test_utils::TestDatabase; use camino::Utf8Path; use camino_tempfile::Utf8TempDir; use nexus_db_model::SCHEMA_VERSION; - use nexus_test_utils::db as test_db; use omicron_test_utils::dev; - use std::sync::Arc; // Confirms that calling the internal "ensure_schema" function can succeed // when the database is already at that version. #[tokio::test] async fn ensure_schema_is_current_version() { let logctx = dev::test_setup_log("ensure_schema_is_current_version"); - let mut crdb = test_db::test_setup_database(&logctx.log).await; - - let cfg = db::Config { url: crdb.pg_config().clone() }; - let pool = Arc::new(db::Pool::new_single_host(&logctx.log, &cfg)); - let datastore = - Arc::new(DataStore::new(&logctx.log, pool, None).await.unwrap()); + let db = TestDatabase::new_with_raw_datastore(&logctx.log).await; + let datastore = db.datastore(); datastore .ensure_schema(&logctx.log, SCHEMA_VERSION, None) .await .expect("Failed to ensure schema"); - datastore.terminate().await; - crdb.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -557,10 +551,8 @@ mod test { let logctx = dev::test_setup_log("concurrent_nexus_instances_only_move_forward"); let log = &logctx.log; - let mut crdb = test_db::test_setup_database(&logctx.log).await; - - let cfg = db::Config { url: crdb.pg_config().clone() }; - let pool = Arc::new(db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // Mimic the layout of "schema/crdb". @@ -660,8 +652,7 @@ mod test { .collect::, _>>() .expect("Failed to create datastore"); - pool.terminate().await; - crdb.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -670,10 +661,8 @@ mod test { let logctx = dev::test_setup_log("schema_version_subcomponents_save_progress"); let log = &logctx.log; - let mut crdb = test_db::test_setup_database(&logctx.log).await; - - let cfg = db::Config { url: crdb.pg_config().clone() }; - let pool = Arc::new(db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // Mimic the layout of "schema/crdb". @@ -781,8 +770,7 @@ mod test { .expect("Failed to get data"); assert_eq!(data, "abcd"); - datastore.terminate().await; - crdb.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/deployment.rs b/nexus/db-queries/src/db/datastore/deployment.rs index af3c1d07cd..4398c4b13f 100644 --- a/nexus/db-queries/src/db/datastore/deployment.rs +++ b/nexus/db-queries/src/db/datastore/deployment.rs @@ -1730,7 +1730,7 @@ impl RunQueryDsl for InsertTargetQuery {} mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::raw_query_builder::QueryBuilder; use nexus_inventory::now_db_precision; use nexus_inventory::CollectionBuilder; @@ -1738,7 +1738,6 @@ mod tests { use nexus_reconfigurator_planning::blueprint_builder::Ensure; use nexus_reconfigurator_planning::blueprint_builder::EnsureMultiple; use nexus_reconfigurator_planning::example::example; - use nexus_test_utils::db::test_setup_database; use nexus_types::deployment::blueprint_zone_type; use nexus_types::deployment::BlueprintZoneConfig; use nexus_types::deployment::BlueprintZoneDisposition; @@ -1906,8 +1905,8 @@ mod tests { async fn test_empty_blueprint() { // Setup let logctx = dev::test_setup_log("test_empty_blueprint"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create an empty blueprint from it let blueprint1 = BlueprintBuilder::build_empty_with_sleds( @@ -1956,8 +1955,7 @@ mod tests { // on other tests to check blueprint deletion. // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1966,8 +1964,8 @@ mod tests { const TEST_NAME: &str = "test_representative_blueprint"; // Setup let logctx = dev::test_setup_log(TEST_NAME); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a cohesive representative collection/policy/blueprint let (collection, planning_input, blueprint1) = @@ -2192,8 +2190,7 @@ mod tests { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2201,8 +2198,8 @@ mod tests { async fn test_set_target() { // Setup let logctx = dev::test_setup_log("test_set_target"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Trying to insert a target that doesn't reference a blueprint should // fail with a relevant error message. @@ -2379,8 +2376,7 @@ mod tests { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2388,8 +2384,8 @@ mod tests { async fn test_set_target_enabled() { // Setup let logctx = dev::test_setup_log("test_set_target_enabled"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create an initial empty collection let collection = CollectionBuilder::new("test").build(); @@ -2493,8 +2489,7 @@ mod tests { } // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2584,8 +2579,8 @@ mod tests { let logctx = dev::test_setup_log( "test_ensure_external_networking_works_with_good_target", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let blueprint = create_blueprint_with_external_ip(&datastore, &opctx).await; @@ -2611,8 +2606,7 @@ mod tests { .expect("Should be able to allocate external network resources"); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2622,8 +2616,8 @@ mod tests { let logctx = dev::test_setup_log( "test_ensure_external_networking_bails_on_bad_target", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create an initial empty collection let collection = CollectionBuilder::new("test").build(); @@ -2817,8 +2811,7 @@ mod tests { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } diff --git a/nexus/db-queries/src/db/datastore/deployment/external_networking.rs b/nexus/db-queries/src/db/datastore/deployment/external_networking.rs index 302a79160d..0f13050b98 100644 --- a/nexus/db-queries/src/db/datastore/deployment/external_networking.rs +++ b/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::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::queries::ALLOW_FULL_TABLE_SCAN_SQL; use anyhow::Context as _; use async_bb8_diesel::AsyncSimpleConnection; @@ -417,7 +417,6 @@ mod tests { use nexus_config::NUM_INITIAL_RESERVED_IP_ADDRESSES; use nexus_db_model::SqlU16; use nexus_sled_agent_shared::inventory::OmicronZoneDataset; - use nexus_test_utils::db::test_setup_database; use nexus_types::deployment::blueprint_zone_type; use nexus_types::deployment::BlueprintZoneConfig; use nexus_types::deployment::BlueprintZoneDisposition; @@ -876,8 +875,8 @@ mod tests { // Set up. usdt::register_probes().unwrap(); let logctx = dev::test_setup_log("test_service_ip_list"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Generate the test values we care about. let mut harness = Harness::new(); @@ -1128,8 +1127,7 @@ mod tests { } // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1138,8 +1136,8 @@ mod tests { // Set up. usdt::register_probes().unwrap(); let logctx = dev::test_setup_log("test_service_ip_list"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Generate the test values we care about. let harness = Harness::new(); @@ -1196,8 +1194,7 @@ mod tests { harness.assert_nics_are_deleted_in_datastore(&datastore).await; // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/disk.rs b/nexus/db-queries/src/db/datastore/disk.rs index d369cd17be..76f4055373 100644 --- a/nexus/db-queries/src/db/datastore/disk.rs +++ b/nexus/db-queries/src/db/datastore/disk.rs @@ -843,8 +843,7 @@ impl DataStore { mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use nexus_types::external_api::params; use omicron_common::api::external; use omicron_test_utils::dev; @@ -854,8 +853,8 @@ mod tests { let logctx = dev::test_setup_log("test_undelete_disk_set_faulted_idempotent"); let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let (opctx, db_datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&log).await; + let (opctx, db_datastore) = (db.opctx(), db.datastore()); let silo_id = opctx.authn.actor().unwrap().silo_id().unwrap(); @@ -979,8 +978,7 @@ mod tests { ); } - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/dns.rs b/nexus/db-queries/src/db/datastore/dns.rs index e86c61b8b9..9279933e47 100644 --- a/nexus/db-queries/src/db/datastore/dns.rs +++ b/nexus/db-queries/src/db/datastore/dns.rs @@ -729,7 +729,7 @@ impl DataStoreDnsTest for DataStore { #[cfg(test)] mod test { - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::DnsVersionUpdateBuilder; use crate::db::DataStore; use crate::db::TransactionError; @@ -744,7 +744,6 @@ mod test { use nexus_db_model::DnsZone; use nexus_db_model::Generation; use nexus_db_model::InitialDnsGroup; - use nexus_test_utils::db::test_setup_database; use nexus_types::internal_api::params::DnsRecord; use nexus_types::internal_api::params::Srv; use omicron_common::api::external::Error; @@ -758,8 +757,8 @@ mod test { #[tokio::test] async fn test_read_dns_config_uninitialized() { let logctx = dev::test_setup_log("test_read_dns_config_uninitialized"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // If we attempt to load the config when literally nothing related to // DNS has been initialized, we will get an InternalError because we @@ -834,8 +833,7 @@ mod test { version for DNS group External, found 0" ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -843,8 +841,8 @@ mod test { #[tokio::test] async fn test_read_dns_config_basic() { let logctx = dev::test_setup_log("test_read_dns_config_basic"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create exactly one zone with no names in it. // This will not show up in the read config. @@ -941,8 +939,7 @@ mod test { .expect("failed to read DNS config with batch size 1"); assert_eq!(dns_config_batch_1, dns_config); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -950,8 +947,8 @@ mod test { #[tokio::test] async fn test_read_dns_config_complex() { let logctx = dev::test_setup_log("test_read_dns_config_complex"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let batch_size = NonZeroU32::new(10).unwrap(); let now = Utc::now(); let log = &logctx.log; @@ -1312,8 +1309,7 @@ mod test { HashMap::from([("n1".to_string(), records_r2.clone())]) ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1321,8 +1317,8 @@ mod test { #[tokio::test] async fn test_dns_uniqueness() { let logctx = dev::test_setup_log("test_dns_uniqueness"); - let mut db = test_setup_database(&logctx.log).await; - let (_opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let now = Utc::now(); // There cannot be two DNS zones in the same group with the same name. @@ -1418,8 +1414,7 @@ mod test { .contains("duplicate key value violates unique constraint")); } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1495,8 +1490,8 @@ mod test { #[tokio::test] async fn test_dns_update_incremental() { let logctx = dev::test_setup_log("test_dns_update_incremental"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let now = Utc::now(); // Create three DNS zones for testing: @@ -1866,16 +1861,15 @@ mod test { assert_eq!(dns_config.zones[1].zone_name, "oxide2.test"); assert_eq!(dns_config.zones[0].records, dns_config.zones[1].records,); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_dns_update_from_version() { let logctx = dev::test_setup_log("test_dns_update_from_version"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // The guts of `dns_update_from_version()` are shared with // `dns_update_incremental()`. The main cases worth testing here are @@ -1980,8 +1974,7 @@ mod test { assert!(!records.contains_key("krabappel")); assert!(records.contains_key("hoover")); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/external_ip.rs b/nexus/db-queries/src/db/datastore/external_ip.rs index 2972c78e56..ef6716a43b 100644 --- a/nexus/db-queries/src/db/datastore/external_ip.rs +++ b/nexus/db-queries/src/db/datastore/external_ip.rs @@ -1132,8 +1132,7 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use nexus_types::deployment::OmicronZoneExternalFloatingIp; use nexus_types::deployment::OmicronZoneExternalSnatIp; use nexus_types::external_api::shared::IpRange; @@ -1164,8 +1163,8 @@ mod tests { async fn test_service_ip_list() { usdt::register_probes().unwrap(); let logctx = dev::test_setup_log("test_service_ip_list"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // No IPs, to start let ips = read_all_service_ips(&datastore, &opctx).await; @@ -1246,8 +1245,7 @@ mod tests { let ips = read_all_service_ips(&datastore, &opctx).await; assert_eq!(ips, external_ips); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/instance.rs b/nexus/db-queries/src/db/datastore/instance.rs index c2f426ceb1..0698883891 100644 --- a/nexus/db-queries/src/db/datastore/instance.rs +++ b/nexus/db-queries/src/db/datastore/instance.rs @@ -1943,15 +1943,14 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::sled; - use crate::db::datastore::test_utils::datastore_test; use crate::db::lookup::LookupPath; use crate::db::pagination::Paginator; use nexus_db_model::InstanceState; use nexus_db_model::Project; use nexus_db_model::VmmRuntimeState; use nexus_db_model::VmmState; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params; use nexus_types::identity::Asset; use nexus_types::silo::DEFAULT_SILO_ID; @@ -2034,8 +2033,8 @@ mod tests { async fn test_instance_updater_acquires_lock() { // Setup let logctx = dev::test_setup_log("test_instance_updater_acquires_lock"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let saga1 = Uuid::new_v4(); let saga2 = Uuid::new_v4(); let (authz_project, _) = create_test_project(&datastore, &opctx).await; @@ -2108,8 +2107,7 @@ mod tests { assert!(unlocked, "instance must actually be unlocked"); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2118,8 +2116,8 @@ mod tests { // Setup let logctx = dev::test_setup_log("test_instance_updater_lock_is_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2173,8 +2171,7 @@ mod tests { assert!(!unlocked, "instance should already have been unlocked"); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2184,8 +2181,8 @@ mod tests { let logctx = dev::test_setup_log( "test_instance_updater_cant_unlock_someone_elses_instance_", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2267,8 +2264,7 @@ mod tests { assert!(!unlocked); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2277,8 +2273,8 @@ mod tests { // Setup let logctx = dev::test_setup_log("test_unlocking_a_deleted_instance_is_okay"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2327,8 +2323,7 @@ mod tests { .expect("instance should unlock"); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2337,8 +2332,8 @@ mod tests { // Setup let logctx = dev::test_setup_log("test_instance_commit_update_is_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2426,8 +2421,7 @@ mod tests { assert_eq!(instance.runtime().r#gen, new_runtime.r#gen); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2437,8 +2431,8 @@ mod tests { let logctx = dev::test_setup_log( "test_instance_update_invalidated_while_locked", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2517,8 +2511,7 @@ mod tests { assert_eq!(instance.runtime().nexus_state, new_runtime.nexus_state); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2526,8 +2519,8 @@ mod tests { async fn test_instance_fetch_all() { // Setup let logctx = dev::test_setup_log("test_instance_fetch_all"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2698,8 +2691,7 @@ mod tests { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2707,8 +2699,8 @@ mod tests { async fn test_instance_set_migration_ids() { // Setup let logctx = dev::test_setup_log("test_instance_set_migration_ids"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let authz_instance = create_test_instance( &datastore, @@ -2966,8 +2958,7 @@ mod tests { assert_eq!(instance.runtime().dst_propolis_id, Some(vmm3.id)); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2977,8 +2968,8 @@ mod tests { // Setup let logctx = dev::test_setup_log("test_instance_and_vmm_list_by_sled_agent"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, _) = create_test_project(&datastore, &opctx).await; let mut expected_instances = BTreeSet::new(); @@ -3104,8 +3095,7 @@ mod tests { assert_eq!(expected_instances, found_instances); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/inventory.rs b/nexus/db-queries/src/db/datastore/inventory.rs index 7dba515f05..4a2ab216a2 100644 --- a/nexus/db-queries/src/db/datastore/inventory.rs +++ b/nexus/db-queries/src/db/datastore/inventory.rs @@ -2444,7 +2444,7 @@ impl DataStoreInventoryTest for DataStore { #[cfg(test)] mod test { use crate::db::datastore::inventory::DataStoreInventoryTest; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::DataStoreConnection; use crate::db::raw_query_builder::{QueryBuilder, TrustedStr}; use crate::db::schema; @@ -2457,7 +2457,6 @@ mod test { use gateway_client::types::SpType; use nexus_inventory::examples::representative; use nexus_inventory::examples::Representative; - use nexus_test_utils::db::test_setup_database; use nexus_test_utils::db::ALLOW_FULL_TABLE_SCAN_SQL; use nexus_types::inventory::BaseboardId; use nexus_types::inventory::CabooseWhich; @@ -2511,8 +2510,8 @@ mod test { #[tokio::test] async fn test_find_hw_baseboard_id_missing_returns_not_found() { let logctx = dev::test_setup_log("inventory_insert"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let baseboard_id = BaseboardId { serial_number: "some-serial".into(), part_number: "some-part".into(), @@ -2522,8 +2521,7 @@ mod test { .await .unwrap_err(); assert!(matches!(err, Error::ObjectNotFound { .. })); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2533,8 +2531,8 @@ mod test { async fn test_inventory_insert() { // Setup let logctx = dev::test_setup_log("inventory_insert"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create an empty collection and write it to the database. let builder = nexus_inventory::CollectionBuilder::new("test"); @@ -3016,8 +3014,7 @@ mod test { assert_ne!(coll_counts.rot_pages, 0); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3100,8 +3097,8 @@ mod test { async fn test_inventory_deletion() { // Setup let logctx = dev::test_setup_log("inventory_deletion"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a representative collection and write it to the database. let Representative { builder, .. } = representative(); @@ -3138,8 +3135,7 @@ mod test { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3147,8 +3143,8 @@ mod test { async fn test_representative_collection_populates_database() { // Setup let logctx = dev::test_setup_log("inventory_deletion"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a representative collection and write it to the database. let Representative { builder, .. } = representative(); @@ -3164,8 +3160,7 @@ mod test { .expect("All inv_... tables should be populated by representative collection"); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/ip_pool.rs b/nexus/db-queries/src/db/datastore/ip_pool.rs index eb7d2aa9c3..9ea8f7b088 100644 --- a/nexus/db-queries/src/db/datastore/ip_pool.rs +++ b/nexus/db-queries/src/db/datastore/ip_pool.rs @@ -1121,12 +1121,11 @@ mod test { use std::num::NonZeroU32; use crate::authz; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::model::{ IpPool, IpPoolResource, IpPoolResourceType, Project, }; use assert_matches::assert_matches; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params; use nexus_types::identity::Resource; use omicron_common::address::{IpRange, Ipv4Range, Ipv6Range}; @@ -1139,8 +1138,8 @@ mod test { #[tokio::test] async fn test_default_ip_pools() { let logctx = dev::test_setup_log("test_default_ip_pools"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // we start out with no default pool, so we expect not found let error = datastore.ip_pools_fetch_default(&opctx).await.unwrap_err(); @@ -1298,16 +1297,15 @@ mod test { .expect("Should list silo IP pools"); assert_eq!(silo_pools.len(), 0); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_internal_ip_pool() { let logctx = dev::test_setup_log("test_internal_ip_pool"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // confirm internal pool appears as internal let (authz_pool, _pool) = @@ -1353,16 +1351,15 @@ mod test { datastore.ip_pool_is_internal(&opctx, &authz_other_pool).await; assert_eq!(is_internal, Ok(false)); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_ip_pool_utilization() { let logctx = dev::test_setup_log("test_ip_utilization"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let authz_silo = opctx.authn.silo_required().unwrap(); let project = Project::new( @@ -1505,8 +1502,7 @@ mod test { assert_eq!(max_ips.ipv4, 5); assert_eq!(max_ips.ipv6, 1208925819614629174706166); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/ipv4_nat_entry.rs b/nexus/db-queries/src/db/datastore/ipv4_nat_entry.rs index 6c0676d4ee..80794f193a 100644 --- a/nexus/db-queries/src/db/datastore/ipv4_nat_entry.rs +++ b/nexus/db-queries/src/db/datastore/ipv4_nat_entry.rs @@ -379,10 +379,9 @@ fn ipv4_nat_next_version() -> diesel::expression::SqlLiteral { mod test { use std::{net::Ipv4Addr, str::FromStr}; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use chrono::Utc; use nexus_db_model::{Ipv4NatEntry, Ipv4NatValues, MacAddr, Vni}; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external; use omicron_test_utils::dev; use rand::seq::IteratorRandom; @@ -391,8 +390,8 @@ mod test { #[tokio::test] async fn nat_version_tracking() { let logctx = dev::test_setup_log("test_nat_version_tracking"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // We should not have any NAT entries at this moment let initial_state = @@ -538,8 +537,7 @@ mod test { 3 ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -549,8 +547,8 @@ mod test { /// of properties. async fn table_allows_unique_active_multiple_deleted() { let logctx = dev::test_setup_log("test_nat_version_tracking"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // We should not have any NAT entries at this moment let initial_state = @@ -683,8 +681,7 @@ mod test { 4 ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -692,8 +689,8 @@ mod test { #[tokio::test] async fn ipv4_nat_sync_service_zones() { let logctx = dev::test_setup_log("ipv4_nat_sync_service_zones"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // We should not have any NAT entries at this moment let initial_state = @@ -806,8 +803,7 @@ mod test { && entry.version_removed.is_none() })); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -815,8 +811,8 @@ mod test { #[tokio::test] async fn ipv4_nat_changeset() { let logctx = dev::test_setup_log("test_nat_version_tracking"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // We should not have any NAT entries at this moment let initial_state = @@ -956,8 +952,7 @@ mod test { // did we see everything? assert_eq!(total_changes, db_records.len()); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/migration.rs b/nexus/db-queries/src/db/datastore/migration.rs index a1155f8bb7..e1d8c070e7 100644 --- a/nexus/db-queries/src/db/datastore/migration.rs +++ b/nexus/db-queries/src/db/datastore/migration.rs @@ -178,11 +178,10 @@ impl DataStore { mod tests { use super::*; use crate::authz; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::lookup::LookupPath; use crate::db::model::Instance; use nexus_db_model::Project; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params; use nexus_types::silo::DEFAULT_SILO_ID; use omicron_common::api::external::ByteCount; @@ -277,8 +276,8 @@ mod tests { async fn test_migration_query_by_instance() { // Setup let logctx = dev::test_setup_log("test_migration_query_by_instance"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let authz_instance = create_test_instance(&datastore, &opctx).await; let instance_id = InstanceUuid::from_untyped_uuid(authz_instance.id()); @@ -356,8 +355,7 @@ mod tests { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } diff --git a/nexus/db-queries/src/db/datastore/mod.rs b/nexus/db-queries/src/db/datastore/mod.rs index 6516d35aca..d2d376b1f8 100644 --- a/nexus/db-queries/src/db/datastore/mod.rs +++ b/nexus/db-queries/src/db/datastore/mod.rs @@ -445,7 +445,7 @@ mod test { use crate::authn; use crate::authn::SiloAuthnPolicy; use crate::authz; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::test_utils::{ IneligibleSledKind, IneligibleSleds, }; @@ -466,7 +466,6 @@ mod test { use nexus_db_fixed_data::silo::DEFAULT_SILO; use nexus_db_model::IpAttachState; use nexus_db_model::{to_db_typed_uuid, Generation}; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params; use nexus_types::silo::DEFAULT_SILO_ID; use omicron_common::api::external::{ @@ -508,8 +507,8 @@ mod test { #[tokio::test] async fn test_project_creation() { let logctx = dev::test_setup_log("test_project_creation"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let authz_silo = opctx.authn.silo_required().unwrap(); @@ -538,16 +537,15 @@ mod test { .unwrap(); assert!(silo_after_project_create.rcgen > silo.rcgen); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_session_methods() { let logctx = dev::test_setup_log("test_session_methods"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let authn_opctx = OpContext::for_background( logctx.log.new(o!("component" => "TestExternalAuthn")), Arc::new(authz::Authz::new(&logctx.log)), @@ -674,8 +672,7 @@ mod test { datastore.session_hard_delete(&opctx, &authz_session).await; assert_eq!(delete_again, Ok(())); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1005,8 +1002,8 @@ mod test { /// pool IDs should not matter. async fn test_region_allocation_strat_random() { let logctx = dev::test_setup_log("test_region_allocation_strat_random"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let test_datasets = TestDatasets::create( &opctx, datastore.clone(), @@ -1083,8 +1080,7 @@ mod test { } } - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1096,8 +1092,8 @@ mod test { let logctx = dev::test_setup_log( "test_region_allocation_strat_random_with_distinct_sleds", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a rack with enough sleds for a successful allocation when we // require 3 distinct eligible sleds. @@ -1173,8 +1169,7 @@ mod test { } } - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1185,8 +1180,8 @@ mod test { let logctx = dev::test_setup_log( "test_region_allocation_strat_random_with_distinct_sleds_fails", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a rack without enough sleds for a successful allocation when // we require 3 distinct provisionable sleds. @@ -1230,8 +1225,7 @@ mod test { assert!(matches!(err, Error::InsufficientCapacity { .. })); } - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1239,8 +1233,8 @@ mod test { async fn test_region_allocation_is_idempotent() { let logctx = dev::test_setup_log("test_region_allocation_is_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); TestDatasets::create( &opctx, datastore.clone(), @@ -1297,8 +1291,7 @@ mod test { assert_eq!(dataset_and_regions1[i], dataset_and_regions2[i],); } - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1307,8 +1300,8 @@ mod test { let logctx = dev::test_setup_log( "test_region_allocation_only_operates_on_zpools_in_inventory", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a sled... let sled_id = create_test_sled(&datastore).await; @@ -1398,8 +1391,7 @@ mod test { .await .expect("Allocation should have worked after adding zpools to inventory"); - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1407,8 +1399,8 @@ mod test { async fn test_region_allocation_not_enough_zpools() { let logctx = dev::test_setup_log("test_region_allocation_not_enough_zpools"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a sled... let sled_id = create_test_sled(&datastore).await; @@ -1484,8 +1476,7 @@ mod test { assert!(matches!(err, Error::InsufficientCapacity { .. })); - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1494,8 +1485,8 @@ mod test { let logctx = dev::test_setup_log( "test_region_allocation_only_considers_disks_in_service", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a sled... let sled_id = create_test_sled(&datastore).await; @@ -1603,8 +1594,7 @@ mod test { } } - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1612,8 +1602,8 @@ mod test { async fn test_region_allocation_out_of_space_fails() { let logctx = dev::test_setup_log("test_region_allocation_out_of_space_fails"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); TestDatasets::create( &opctx, @@ -1638,8 +1628,7 @@ mod test { .await .is_err()); - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1650,11 +1639,8 @@ mod test { use omicron_common::api::external; let logctx = dev::test_setup_log("test_queries_do_not_require_full_table_scan"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); - let datastore = - DataStore::new(&logctx.log, Arc::new(pool), None).await.unwrap(); + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let conn = datastore.pool_connection_for_tests().await.unwrap(); let explanation = DataStore::get_allocated_regions_query(Uuid::nil()) .explain_async(&conn) @@ -1685,8 +1671,7 @@ mod test { explanation, ); - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -1696,15 +1681,8 @@ mod test { use std::net::Ipv6Addr; let logctx = dev::test_setup_log("test_sled_ipv6_address_allocation"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = Arc::new(db::Pool::new_single_host(&logctx.log, &cfg)); - let datastore = - Arc::new(DataStore::new(&logctx.log, pool, None).await.unwrap()); - let opctx = OpContext::for_tests( - logctx.log.new(o!()), - Arc::clone(&datastore) as Arc, - ); + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let rack_id = Uuid::new_v4(); let addr1 = "[fd00:1de::1]:12345".parse().unwrap(); @@ -1744,16 +1722,15 @@ mod test { let expected_ip = Ipv6Addr::new(0xfd00, 0x1df, 0, 0, 0, 0, 1, 0); assert_eq!(ip, expected_ip); - datastore.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_ssh_keys() { let logctx = dev::test_setup_log("test_ssh_keys"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a new Silo user so that we can lookup their keys. let authz_silo = authz::Silo::new( @@ -1829,16 +1806,15 @@ mod test { datastore.ssh_key_delete(&opctx, &authz_ssh_key).await.unwrap(); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_rack_initialize_is_idempotent() { let logctx = dev::test_setup_log("test_rack_initialize_is_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a Rack, insert it into the DB. let rack = Rack::new(Uuid::new_v4()); @@ -1871,16 +1847,15 @@ mod test { .unwrap(); assert!(result.initialized); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_table_scan() { let logctx = dev::test_setup_log("test_table_scan"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let error = datastore.test_try_table_scan(&opctx).await; println!("error from attempted table scan: {:#}", error); @@ -1898,8 +1873,7 @@ mod test { } // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1911,8 +1885,8 @@ mod test { let logctx = dev::test_setup_log( "test_deallocate_external_ip_by_instance_id_is_idempotent", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let conn = datastore.pool_connection_for_tests().await.unwrap(); // Create a few records. @@ -1966,8 +1940,7 @@ mod test { .expect("Failed to delete instance external IPs"); assert_eq!(count, 0, "Expected to delete zero IPs for the instance"); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1978,8 +1951,8 @@ mod test { let logctx = dev::test_setup_log("test_deallocate_external_ip_is_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let conn = datastore.pool_connection_for_tests().await.unwrap(); // Create a record. @@ -2033,8 +2006,7 @@ mod test { .await .is_err()); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2047,8 +2019,8 @@ mod test { use diesel::result::Error::DatabaseError; let logctx = dev::test_setup_log("test_external_ip_check_constraints"); - let mut db = test_setup_database(&logctx.log).await; - let (_opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let conn = datastore.pool_connection_for_tests().await.unwrap(); let now = Utc::now(); @@ -2282,8 +2254,7 @@ mod test { } } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/network_interface.rs b/nexus/db-queries/src/db/datastore/network_interface.rs index a652f6a5a5..b7f0622609 100644 --- a/nexus/db-queries/src/db/datastore/network_interface.rs +++ b/nexus/db-queries/src/db/datastore/network_interface.rs @@ -892,10 +892,9 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use nexus_config::NUM_INITIAL_RESERVED_IP_ADDRESSES; use nexus_db_fixed_data::vpc_subnet::NEXUS_VPC_SUBNET; - use nexus_test_utils::db::test_setup_database; use omicron_common::address::NEXUS_OPTE_IPV4_SUBNET; use omicron_test_utils::dev; use std::collections::BTreeSet; @@ -924,8 +923,8 @@ mod tests { usdt::register_probes().unwrap(); let logctx = dev::test_setup_log("test_service_network_interfaces_list"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // No IPs, to start let nics = read_all_service_nics(&datastore, &opctx).await; @@ -991,8 +990,7 @@ mod tests { let nics = read_all_service_nics(&datastore, &opctx).await; assert_eq!(nics, service_nics); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/oximeter.rs b/nexus/db-queries/src/db/datastore/oximeter.rs index d8632b096a..be5ddb91bb 100644 --- a/nexus/db-queries/src/db/datastore/oximeter.rs +++ b/nexus/db-queries/src/db/datastore/oximeter.rs @@ -292,8 +292,7 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; - use db::datastore::pub_test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use nexus_types::internal_api::params; use omicron_common::api::internal::nexus; use omicron_test_utils::dev; @@ -345,9 +344,8 @@ mod tests { async fn test_oximeter_expunge() { // Setup let logctx = dev::test_setup_log("test_oximeter_expunge"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Insert a few Oximeter collectors. let mut collector_ids = @@ -447,8 +445,7 @@ mod tests { assert_eq!(expunged1a, expunged1b); // Cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -458,9 +455,8 @@ mod tests { let logctx = dev::test_setup_log( "test_producer_endpoint_reassigns_if_oximeter_expunged", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Insert an Oximeter collector. let oximeter1_id = Uuid::new_v4(); @@ -575,8 +571,7 @@ mod tests { } // Cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -586,9 +581,8 @@ mod tests { let logctx = dev::test_setup_log( "test_producer_endpoint_upsert_rejects_expunged_oximeters", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Insert a few Oximeter collectors. let collector_ids = (0..4).map(|_| Uuid::new_v4()).collect::>(); @@ -686,8 +680,7 @@ mod tests { ); // Cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -695,9 +688,8 @@ mod tests { async fn test_oximeter_reassigns_randomly() { // Setup let logctx = dev::test_setup_log("test_oximeter_reassigns_randomly"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Insert a few Oximeter collectors. let collector_ids = (0..4).map(|_| Uuid::new_v4()).collect::>(); @@ -791,8 +783,7 @@ mod tests { assert_eq!(producer_counts[1..].iter().sum::(), 1000); // Cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -802,9 +793,8 @@ mod tests { let logctx = dev::test_setup_log( "test_oximeter_reassign_fails_if_no_collectors", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Insert a few Oximeter collectors. let collector_ids = (0..4).map(|_| Uuid::new_v4()).collect::>(); @@ -899,8 +889,7 @@ mod tests { assert_eq!(nproducers, 100); // Cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -908,9 +897,8 @@ mod tests { async fn test_producers_list_expired() { // Setup let logctx = dev::test_setup_log("test_producers_list_expired"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Insert an Oximeter collector let collector_info = OximeterInfo::new(¶ms::OximeterInfo { @@ -977,8 +965,7 @@ mod tests { .await; assert_eq!(expired_producers.as_slice(), &[]); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/physical_disk.rs b/nexus/db-queries/src/db/datastore/physical_disk.rs index 3057954e13..73aa837af8 100644 --- a/nexus/db-queries/src/db/datastore/physical_disk.rs +++ b/nexus/db-queries/src/db/datastore/physical_disk.rs @@ -321,10 +321,10 @@ 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::datastore::test_utils::datastore_test; use crate::db::lookup::LookupPath; use crate::db::model::{PhysicalDiskKind, Sled, SledUpdate}; use dropshot::PaginationOrder; @@ -332,7 +332,6 @@ mod test { use nexus_sled_agent_shared::inventory::{ Baseboard, Inventory, InventoryDisk, OmicronZonesConfig, SledRole, }; - use nexus_test_utils::db::test_setup_database; use nexus_types::identity::Asset; use omicron_common::api::external::ByteCount; use omicron_common::disk::{DiskIdentity, DiskVariant}; @@ -372,8 +371,8 @@ mod test { async fn physical_disk_insert_same_uuid_collides() { let logctx = dev::test_setup_log("physical_disk_insert_same_uuid_collides"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore).await; let sled_id = sled.id(); @@ -405,8 +404,7 @@ mod test { "{err}" ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -414,8 +412,8 @@ mod test { async fn physical_disk_insert_different_disks() { let logctx = dev::test_setup_log("physical_disk_insert_different_disks"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore).await; let sled_id = sled.id(); @@ -455,16 +453,15 @@ mod test { .expect("Failed to list physical disks"); assert_eq!(disks.len(), 2); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn physical_disk_deletion_idempotency() { let logctx = dev::test_setup_log("physical_disk_deletion_idempotency"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore).await; @@ -506,8 +503,7 @@ mod test { .await .expect("Failed to delete disk"); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -521,8 +517,8 @@ mod test { let logctx = dev::test_setup_log( "physical_disk_insert_delete_reupsert_new_sled", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled_a = create_test_sled(&datastore).await; let sled_b = create_test_sled(&datastore).await; @@ -595,8 +591,7 @@ mod test { .expect("Failed to list physical disks"); assert_eq!(disks.len(), 1); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -611,8 +606,8 @@ mod test { async fn physical_disk_insert_reupsert_new_sled() { let logctx = dev::test_setup_log("physical_disk_insert_reupsert_new_sled"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled_a = create_test_sled(&datastore).await; let sled_b = create_test_sled(&datastore).await; @@ -674,8 +669,7 @@ mod test { .expect("Failed to list physical disks"); assert_eq!(disks.len(), 1); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -781,8 +775,8 @@ mod test { async fn physical_disk_cannot_insert_to_expunged_sled() { let logctx = dev::test_setup_log("physical_disk_cannot_insert_to_expunged_sled"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore).await; @@ -819,16 +813,15 @@ mod test { "Expected string: {expected} within actual error: {actual}", ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn physical_disk_uninitialized_list() { let logctx = dev::test_setup_log("physical_disk_uninitialized_list"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled_a = create_test_sled(&datastore).await; let sled_b = create_test_sled(&datastore).await; @@ -1006,8 +999,7 @@ mod test { .expect("Failed to list uninitialized disks"); assert_eq!(uninitialized_disks.len(), 0); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/pub_test_utils.rs b/nexus/db-queries/src/db/datastore/pub_test_utils.rs index bcf6a6c80f..233113ea83 100644 --- a/nexus/db-queries/src/db/datastore/pub_test_utils.rs +++ b/nexus/db-queries/src/db/datastore/pub_test_utils.rs @@ -12,32 +12,132 @@ use crate::authz; use crate::context::OpContext; use crate::db; use crate::db::DataStore; -use dropshot::test_util::LogContext; use omicron_test_utils::dev::db::CockroachInstance; +use slog::Logger; use std::sync::Arc; use uuid::Uuid; +#[cfg(test)] +mod test { + use super::*; + use nexus_test_utils::db::test_setup_database; + + enum TestState { + Pool { pool: Arc }, + RawDatastore { datastore: Arc }, + Datastore { opctx: OpContext, datastore: Arc }, + } + + /// A test database with a pool connected to it. + pub struct TestDatabase { + db: CockroachInstance, + + state: TestState, + } + + impl TestDatabase { + /// Creates a new database for test usage, with a pool. + /// + /// [Self::terminate] should be called before the test finishes. + pub async fn new_with_pool(log: &Logger) -> Self { + let db = test_setup_database(log).await; + let cfg = db::Config { url: db.pg_config().clone() }; + let pool = Arc::new(db::Pool::new_single_host(log, &cfg)); + Self { db, state: TestState::Pool { pool } } + } + + /// Creates a new database for test usage, with a pre-loaded datastore. + /// + /// [Self::terminate] should be called before the test finishes. + pub async fn new_with_datastore(log: &Logger) -> Self { + let db = test_setup_database(log).await; + let (opctx, datastore) = + crate::db::datastore::test_utils::datastore_test(log, &db) + .await; + + Self { db, state: TestState::Datastore { opctx, datastore } } + } + + /// Creates a new database for test usage, with a raw datastore. + /// + /// [Self::terminate] should be called before the test finishes. + pub async fn new_with_raw_datastore(log: &Logger) -> Self { + let db = test_setup_database(log).await; + let cfg = db::Config { url: db.pg_config().clone() }; + let pool = Arc::new(db::Pool::new_single_host(log, &cfg)); + let datastore = + Arc::new(DataStore::new(&log, pool, None).await.unwrap()); + Self { db, state: TestState::RawDatastore { datastore } } + } + + pub fn pool(&self) -> &Arc { + match &self.state { + TestState::Pool { pool } => pool, + TestState::RawDatastore { .. } + | TestState::Datastore { .. } => { + panic!("Wrong test type; try using `TestDatabase::new_with_pool`"); + } + } + } + + pub fn opctx(&self) -> &OpContext { + match &self.state { + TestState::Pool { .. } | TestState::RawDatastore { .. } => { + panic!("Wrong test type; try using `TestDatabase::new_with_datastore`"); + } + TestState::Datastore { opctx, .. } => opctx, + } + } + + pub fn datastore(&self) -> &Arc { + match &self.state { + TestState::Pool { .. } => { + panic!("Wrong test type; try using `TestDatabase::new_with_datastore`"); + } + TestState::RawDatastore { datastore } => datastore, + TestState::Datastore { datastore, .. } => datastore, + } + } + + /// Shuts down both the database and the pool + pub async fn terminate(mut self) { + match self.state { + TestState::Pool { pool } => pool.terminate().await, + TestState::RawDatastore { datastore } => { + datastore.terminate().await + } + TestState::Datastore { datastore, .. } => { + datastore.terminate().await + } + } + self.db.cleanup().await.unwrap(); + } + } +} + +#[cfg(test)] +pub use test::TestDatabase; + /// Constructs a DataStore for use in test suites that has preloaded the /// built-in users, roles, and role assignments that are needed for basic /// operation #[cfg(any(test, feature = "testing"))] pub async fn datastore_test( - logctx: &LogContext, + log: &Logger, db: &CockroachInstance, rack_id: Uuid, ) -> (OpContext, Arc) { use crate::authn; let cfg = db::Config { url: db.pg_config().clone() }; - let pool = Arc::new(db::Pool::new_single_host(&logctx.log, &cfg)); - let datastore = - Arc::new(DataStore::new(&logctx.log, pool, None).await.unwrap()); + let pool = Arc::new(db::Pool::new_single_host(&log, &cfg)); + let datastore = Arc::new(DataStore::new(&log, pool, None).await.unwrap()); // Create an OpContext with the credentials of "db-init" just for the // purpose of loading the built-in users, roles, and assignments. let opctx = OpContext::for_background( - logctx.log.new(o!()), - Arc::new(authz::Authz::new(&logctx.log)), + log.new(o!()), + Arc::new(authz::Authz::new(&log)), authn::Context::internal_db_init(), Arc::clone(&datastore) as Arc, ); @@ -60,7 +160,7 @@ pub async fn datastore_test( // Create an OpContext with the credentials of "test-privileged" for general // testing. let opctx = OpContext::for_tests( - logctx.log.new(o!()), + log.new(o!()), Arc::clone(&datastore) as Arc, ); diff --git a/nexus/db-queries/src/db/datastore/rack.rs b/nexus/db-queries/src/db/datastore/rack.rs index 4dc88cea93..ea0ef57908 100644 --- a/nexus/db-queries/src/db/datastore/rack.rs +++ b/nexus/db-queries/src/db/datastore/rack.rs @@ -997,10 +997,10 @@ 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::datastore::test_utils::datastore_test; use crate::db::datastore::Discoverability; use crate::db::model::ExternalIp; use crate::db::model::IpKind; @@ -1015,7 +1015,6 @@ mod test { SledBuilder, SystemDescription, }; use nexus_sled_agent_shared::inventory::OmicronZoneDataset; - use nexus_test_utils::db::test_setup_database; use nexus_types::deployment::BlueprintZonesConfig; use nexus_types::deployment::CockroachDbPreserveDowngrade; use nexus_types::deployment::{ @@ -1131,8 +1130,8 @@ mod test { #[tokio::test] async fn rack_set_initialized_empty() { let logctx = dev::test_setup_log("rack_set_initialized_empty"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let before = Utc::now(); let rack_init = RackInit::default(); @@ -1233,8 +1232,7 @@ mod test { .unwrap(); assert_eq!(dns_internal, dns_internal2); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1318,8 +1316,8 @@ mod test { async fn rack_set_initialized_with_services() { let test_name = "rack_set_initialized_with_services"; let logctx = dev::test_setup_log(test_name); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled1 = create_test_sled(&datastore, Uuid::new_v4()).await; let sled2 = create_test_sled(&datastore, Uuid::new_v4()).await; @@ -1662,8 +1660,7 @@ mod test { let observed_datasets = get_all_datasets(&datastore).await; assert!(observed_datasets.is_empty()); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1671,8 +1668,8 @@ mod test { async fn rack_set_initialized_with_many_nexus_services() { let test_name = "rack_set_initialized_with_many_nexus_services"; let logctx = dev::test_setup_log(test_name); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore, Uuid::new_v4()).await; @@ -1948,8 +1945,7 @@ mod test { Some(&external_records) ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1958,8 +1954,8 @@ mod test { let test_name = "rack_set_initialized_missing_service_pool_ip_throws_error"; let logctx = dev::test_setup_log(test_name); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore, Uuid::new_v4()).await; @@ -2049,8 +2045,7 @@ mod test { assert!(get_all_datasets(&datastore).await.is_empty()); assert!(get_all_external_ips(&datastore).await.is_empty()); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2058,8 +2053,8 @@ mod test { async fn rack_set_initialized_overlapping_ips_throws_error() { let test_name = "rack_set_initialized_overlapping_ips_throws_error"; let logctx = dev::test_setup_log(test_name); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sled = create_test_sled(&datastore, Uuid::new_v4()).await; @@ -2199,16 +2194,15 @@ mod test { assert!(get_all_datasets(&datastore).await.is_empty()); assert!(get_all_external_ips(&datastore).await.is_empty()); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn rack_sled_subnet_allocations() { let logctx = dev::test_setup_log("rack_sled_subnet_allocations"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let rack_id = Uuid::new_v4(); @@ -2293,16 +2287,15 @@ mod test { allocations.iter().map(|a| a.subnet_octet).collect::>() ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn allocate_sled_underlay_subnet_octets() { let logctx = dev::test_setup_log("rack_sled_subnet_allocations"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let rack_id = Uuid::new_v4(); @@ -2488,8 +2481,7 @@ mod test { next_expected_octet += 1; } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/region_replacement.rs b/nexus/db-queries/src/db/datastore/region_replacement.rs index 922511ad49..508e80a63b 100644 --- a/nexus/db-queries/src/db/datastore/region_replacement.rs +++ b/nexus/db-queries/src/db/datastore/region_replacement.rs @@ -895,15 +895,14 @@ impl DataStore { mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use omicron_test_utils::dev; #[tokio::test] async fn test_one_replacement_per_volume() { let logctx = dev::test_setup_log("test_one_replacement_per_volume"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let region_1_id = Uuid::new_v4(); let region_2_id = Uuid::new_v4(); @@ -921,8 +920,7 @@ mod test { .await .unwrap_err(); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -935,8 +933,8 @@ mod test { let logctx = dev::test_setup_log( "test_replacement_done_in_middle_of_drive_saga", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let region_id = Uuid::new_v4(); let volume_id = Uuid::new_v4(); @@ -1015,8 +1013,7 @@ mod test { ); assert_eq!(actual_request.operating_saga_id, None); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1028,8 +1025,8 @@ mod test { let logctx = dev::test_setup_log( "test_replacement_done_in_middle_of_finish_saga", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let region_id = Uuid::new_v4(); let volume_id = Uuid::new_v4(); @@ -1083,8 +1080,7 @@ mod test { .await .unwrap(); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/region_snapshot_replacement.rs b/nexus/db-queries/src/db/datastore/region_snapshot_replacement.rs index 43424b345e..d9c8a8b258 100644 --- a/nexus/db-queries/src/db/datastore/region_snapshot_replacement.rs +++ b/nexus/db-queries/src/db/datastore/region_snapshot_replacement.rs @@ -1059,16 +1059,15 @@ impl DataStore { mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::model::RegionReplacement; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; #[tokio::test] async fn test_one_replacement_per_volume() { let logctx = dev::test_setup_log("test_one_replacement_per_volume"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let dataset_1_id = Uuid::new_v4(); let region_1_id = Uuid::new_v4(); @@ -1106,8 +1105,7 @@ mod test { .await .unwrap_err(); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1116,8 +1114,8 @@ mod test { let logctx = dev::test_setup_log( "test_one_replacement_per_volume_conflict_with_region", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let dataset_1_id = Uuid::new_v4(); let region_1_id = Uuid::new_v4(); @@ -1147,16 +1145,15 @@ mod test { .await .unwrap_err(); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn count_replacement_steps() { let logctx = dev::test_setup_log("count_replacement_steps"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let dataset_id = Uuid::new_v4(); let region_id = Uuid::new_v4(); @@ -1302,8 +1299,7 @@ mod test { 1, ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1312,8 +1308,8 @@ mod test { let logctx = dev::test_setup_log( "unique_region_snapshot_replacement_step_per_volume", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Ensure that only one non-complete replacement step can be inserted // per volume. @@ -1404,16 +1400,15 @@ mod test { .await .unwrap(); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn region_snapshot_replacement_step_gc() { let logctx = dev::test_setup_log("region_snapshot_replacement_step_gc"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let mut request = RegionSnapshotReplacement::new( Uuid::new_v4(), @@ -1474,8 +1469,7 @@ mod test { .len(), ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1483,8 +1477,8 @@ mod test { async fn region_snapshot_replacement_step_conflict() { let logctx = dev::test_setup_log("region_snapshot_replacement_step_conflict"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Assert that a region snapshot replacement step cannot be created for // a volume that is the "old snapshot volume" for another snapshot @@ -1525,8 +1519,7 @@ mod test { } ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1536,8 +1529,8 @@ mod test { let logctx = dev::test_setup_log( "region_snapshot_replacement_step_conflict_with_region_replacement", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Assert that a region snapshot replacement step cannot be performed on // a volume if region replacement is occurring for that volume. @@ -1559,8 +1552,7 @@ mod test { .await .unwrap_err(); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/saga.rs b/nexus/db-queries/src/db/datastore/saga.rs index d59f5c4fe5..f1f0bd18cc 100644 --- a/nexus/db-queries/src/db/datastore/saga.rs +++ b/nexus/db-queries/src/db/datastore/saga.rs @@ -259,12 +259,11 @@ impl DataStore { #[cfg(test)] mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use async_bb8_diesel::AsyncConnection; use async_bb8_diesel::AsyncSimpleConnection; use db::queries::ALLOW_FULL_TABLE_SCAN_SQL; use nexus_db_model::{SagaNodeEvent, SecId}; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external::Generation; use omicron_test_utils::dev; use rand::seq::SliceRandom; @@ -276,8 +275,8 @@ mod test { async fn test_list_candidate_sagas() { // Test setup let logctx = dev::test_setup_log("test_list_candidate_sagas"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let sec_id = db::SecId(uuid::Uuid::new_v4()); let mut inserted_sagas = (0..SQL_BATCH_SIZE.get() * 2) .map(|_| SagaTestContext::new(sec_id).new_running_db_saga()) @@ -324,8 +323,7 @@ mod test { ); // Test cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -334,8 +332,8 @@ mod test { async fn test_list_unfinished_nodes() { // Test setup let logctx = dev::test_setup_log("test_list_unfinished_nodes"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let node_cx = SagaTestContext::new(SecId(Uuid::new_v4())); // Create a couple batches of saga events @@ -400,8 +398,7 @@ mod test { } // Test cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -410,8 +407,8 @@ mod test { async fn test_list_no_unfinished_nodes() { // Test setup let logctx = dev::test_setup_log("test_list_no_unfinished_nodes"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let saga_id = steno::SagaId(Uuid::new_v4()); // Test that this returns "no nodes" rather than throwing some "not @@ -426,8 +423,7 @@ mod test { assert_eq!(observed_nodes.len(), 0); // Test cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -435,8 +431,8 @@ mod test { async fn test_create_event_idempotent() { // Test setup let logctx = dev::test_setup_log("test_create_event_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (_, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let node_cx = SagaTestContext::new(SecId(Uuid::new_v4())); // Generate a bunch of events. @@ -469,8 +465,7 @@ mod test { } // Test cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -478,8 +473,8 @@ mod test { async fn test_update_state_idempotent() { // Test setup let logctx = dev::test_setup_log("test_create_event_idempotent"); - let mut db = test_setup_database(&logctx.log).await; - let (_, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let node_cx = SagaTestContext::new(SecId(Uuid::new_v4())); // Create a saga in the running state. @@ -522,8 +517,7 @@ mod test { .expect("updating state to Done again"); // Test cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -568,9 +562,8 @@ mod test { async fn test_saga_reassignment() { // Test setup let logctx = dev::test_setup_log("test_saga_reassignment"); - let mut db = test_setup_database(&logctx.log).await; - let (_, datastore) = datastore_test(&logctx, &db).await; - let opctx = OpContext::for_tests(logctx.log.clone(), datastore.clone()); + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Populate the database with a few different sagas: // @@ -712,8 +705,7 @@ mod test { assert_eq!(nreassigned, 0); // Test cleanup - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/sled.rs b/nexus/db-queries/src/db/datastore/sled.rs index 75043a2115..8e37d7ae7f 100644 --- a/nexus/db-queries/src/db/datastore/sled.rs +++ b/nexus/db-queries/src/db/datastore/sled.rs @@ -824,12 +824,12 @@ impl TransitionError { #[cfg(test)] pub(in crate::db::datastore) 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::datastore::test_utils::{ - datastore_test, sled_set_policy, sled_set_state, Expected, - IneligibleSleds, + sled_set_policy, sled_set_state, Expected, IneligibleSleds, }; use crate::db::lookup::LookupPath; use crate::db::model::ByteCount; @@ -841,7 +841,6 @@ pub(in crate::db::datastore) mod test { use nexus_db_model::PhysicalDiskKind; use nexus_db_model::PhysicalDiskPolicy; use nexus_db_model::PhysicalDiskState; - use nexus_test_utils::db::test_setup_database; use nexus_types::identity::Asset; use omicron_common::api::external; use omicron_test_utils::dev; @@ -857,8 +856,8 @@ pub(in crate::db::datastore) mod test { #[tokio::test] async fn upsert_sled_updates_hardware() { let logctx = dev::test_setup_log("upsert_sled_updates_hardware"); - let mut db = test_setup_database(&logctx.log).await; - let (_opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let mut sled_update = test_new_sled_update(); let (observed_sled, _) = @@ -908,8 +907,7 @@ pub(in crate::db::datastore) mod test { ); assert_eq!(observed_sled.reservoir_size, sled_update.reservoir_size); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -918,8 +916,8 @@ pub(in crate::db::datastore) mod test { let logctx = dev::test_setup_log( "upsert_sled_updates_fails_with_stale_sled_agent_gen", ); - let mut db = test_setup_database(&logctx.log).await; - let (_opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let mut sled_update = test_new_sled_update(); let (observed_sled, _) = @@ -973,8 +971,7 @@ pub(in crate::db::datastore) mod test { assert_eq!(observed_sled.reservoir_size, sled_update.reservoir_size); assert_eq!(observed_sled.sled_agent_gen, sled_update.sled_agent_gen); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -982,8 +979,8 @@ pub(in crate::db::datastore) mod test { async fn upsert_sled_doesnt_update_decommissioned() { let logctx = dev::test_setup_log("upsert_sled_doesnt_update_decommissioned"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let mut sled_update = test_new_sled_update(); let (observed_sled, _) = @@ -1052,8 +1049,7 @@ pub(in crate::db::datastore) mod test { "reservoir_size should not have changed" ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1062,8 +1058,8 @@ pub(in crate::db::datastore) mod test { async fn sled_reservation_create_non_provisionable() { let logctx = dev::test_setup_log("sled_reservation_create_non_provisionable"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Define some sleds that resources cannot be provisioned on. let (non_provisionable_sled, _) = @@ -1140,8 +1136,7 @@ pub(in crate::db::datastore) mod test { .unwrap(); } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1168,9 +1163,8 @@ pub(in crate::db::datastore) mod test { async fn test_sled_expungement_also_expunges_disks() { let logctx = dev::test_setup_log("test_sled_expungement_also_expunges_disks"); - let mut db = test_setup_database(&logctx.log).await; - - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Set up a sled to test against. let (sled, _) = @@ -1266,8 +1260,7 @@ pub(in crate::db::datastore) mod test { lookup_physical_disk(&datastore, disk2.id()).await.disk_state ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1275,9 +1268,8 @@ pub(in crate::db::datastore) mod test { async fn test_sled_transitions() { // Test valid and invalid state and policy transitions. let logctx = dev::test_setup_log("test_sled_transitions"); - let mut db = test_setup_database(&logctx.log).await; - - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // This test generates all possible sets of transitions. Below, we list // the before and after predicates for valid transitions. @@ -1391,8 +1383,7 @@ pub(in crate::db::datastore) mod test { .unwrap(); } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1512,8 +1503,8 @@ pub(in crate::db::datastore) mod test { #[tokio::test] async fn sled_list_batch() { let logctx = dev::test_setup_log("sled_list_batch"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let size = usize::try_from(2 * SQL_BATCH_SIZE.get()).unwrap(); let mut new_sleds = Vec::with_capacity(size); @@ -1554,8 +1545,7 @@ pub(in crate::db::datastore) mod test { assert_eq!(expected_ids, found_ids); assert_eq!(found_ids.len(), size); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/switch_port.rs b/nexus/db-queries/src/db/datastore/switch_port.rs index 27a461328a..b332c57798 100644 --- a/nexus/db-queries/src/db/datastore/switch_port.rs +++ b/nexus/db-queries/src/db/datastore/switch_port.rs @@ -1619,9 +1619,8 @@ async fn do_switch_port_settings_delete( #[cfg(test)] mod test { - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::UpdatePrecondition; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params::{ BgpAnnounceSetCreate, BgpConfigCreate, BgpPeerConfig, SwitchPortConfigCreate, SwitchPortGeometry, SwitchPortSettingsCreate, @@ -1637,8 +1636,8 @@ mod test { #[tokio::test] async fn test_bgp_boundary_switches() { let logctx = dev::test_setup_log("test_bgp_boundary_switches"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let rack_id: Uuid = nexus_test_utils::RACK_UUID.parse().expect("parse uuid"); @@ -1738,8 +1737,7 @@ mod test { assert_eq!(uplink_ports.len(), 1); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/test_utils.rs b/nexus/db-queries/src/db/datastore/test_utils.rs index 4678e07f47..75d8833873 100644 --- a/nexus/db-queries/src/db/datastore/test_utils.rs +++ b/nexus/db-queries/src/db/datastore/test_utils.rs @@ -13,7 +13,6 @@ use anyhow::bail; use anyhow::ensure; use anyhow::Context; use anyhow::Result; -use dropshot::test_util::LogContext; use futures::future::try_join_all; use nexus_db_model::SledState; use nexus_types::external_api::views::SledPolicy; @@ -21,16 +20,17 @@ use nexus_types::external_api::views::SledProvisionPolicy; use omicron_test_utils::dev::db::CockroachInstance; use omicron_uuid_kinds::GenericUuid; use omicron_uuid_kinds::SledUuid; +use slog::Logger; use std::sync::Arc; use strum::EnumCount; use uuid::Uuid; pub(crate) async fn datastore_test( - logctx: &LogContext, + log: &Logger, db: &CockroachInstance, ) -> (OpContext, Arc) { let rack_id = Uuid::parse_str(nexus_test_utils::RACK_UUID).unwrap(); - super::pub_test_utils::datastore_test(logctx, db, rack_id).await + super::pub_test_utils::datastore_test(log, db, rack_id).await } /// Denotes a specific way in which a sled is ineligible. diff --git a/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs b/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs index 08dbb13948..a72c032125 100644 --- a/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs +++ b/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs @@ -326,12 +326,11 @@ impl DataStore { mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::lookup::LookupPath; use nexus_db_model::Instance; use nexus_db_model::Project; use nexus_db_model::SiloQuotasUpdate; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params; use nexus_types::silo::DEFAULT_SILO_ID; use omicron_common::api::external::IdentityMetadataCreateParams; @@ -469,8 +468,8 @@ mod test { #[tokio::test] async fn test_instance_create_and_delete() { let logctx = dev::test_setup_log("test_instance_create_and_delete"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let test_data = setup_collections(&datastore, &opctx).await; let ids = test_data.ids(); @@ -531,8 +530,7 @@ mod test { verify_collection_usage(&datastore, &opctx, id, 0, 0, 0).await; } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -540,8 +538,8 @@ mod test { async fn test_instance_create_and_delete_twice() { let logctx = dev::test_setup_log("test_instance_create_and_delete_twice"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let test_data = setup_collections(&datastore, &opctx).await; let ids = test_data.ids(); @@ -645,16 +643,15 @@ mod test { verify_collection_usage(&datastore, &opctx, id, 0, 0, 0).await; } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn test_storage_create_and_delete() { let logctx = dev::test_setup_log("test_storage_create_and_delete"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let test_data = setup_collections(&datastore, &opctx).await; let ids = test_data.ids(); @@ -701,8 +698,7 @@ mod test { verify_collection_usage(&datastore, &opctx, id, 0, 0, 0).await; } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -710,8 +706,8 @@ mod test { async fn test_storage_create_and_delete_twice() { let logctx = dev::test_setup_log("test_storage_create_and_delete_twice"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let test_data = setup_collections(&datastore, &opctx).await; let ids = test_data.ids(); @@ -800,8 +796,7 @@ mod test { verify_collection_usage(&datastore, &opctx, id, 0, 0, 0).await; } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/vmm.rs b/nexus/db-queries/src/db/datastore/vmm.rs index 1d7120a3f1..e578bb1696 100644 --- a/nexus/db-queries/src/db/datastore/vmm.rs +++ b/nexus/db-queries/src/db/datastore/vmm.rs @@ -441,12 +441,11 @@ impl DataStore { mod tests { use super::*; use crate::db; - use crate::db::datastore::test_utils::datastore_test; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::model::Generation; use crate::db::model::Migration; use crate::db::model::VmmRuntimeState; use crate::db::model::VmmState; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::internal::nexus; use omicron_test_utils::dev; use omicron_uuid_kinds::InstanceUuid; @@ -456,8 +455,8 @@ mod tests { // Setup let logctx = dev::test_setup_log("test_vmm_and_migration_update_runtime"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let instance_id = InstanceUuid::from_untyped_uuid(Uuid::new_v4()); let vmm1 = datastore @@ -724,8 +723,7 @@ mod tests { ); // Clean up. - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/volume.rs b/nexus/db-queries/src/db/datastore/volume.rs index 4b659f44b2..85c1401b82 100644 --- a/nexus/db-queries/src/db/datastore/volume.rs +++ b/nexus/db-queries/src/db/datastore/volume.rs @@ -2780,8 +2780,7 @@ impl DataStore { mod tests { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use omicron_test_utils::dev; use sled_agent_client::types::CrucibleOpts; @@ -2792,13 +2791,13 @@ mod tests { let logctx = dev::test_setup_log("test_deserialize_old_crucible_resources"); let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let (_opctx, db_datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&log).await; + let datastore = db.datastore(); // Start with a fake volume, doesn't matter if it's empty let volume_id = Uuid::new_v4(); - let _volume = db_datastore + let _volume = datastore .volume_create(nexus_db_model::Volume::new( volume_id, serde_json::to_string(&VolumeConstructionRequest::Volume { @@ -2819,8 +2818,7 @@ mod tests { { use db::schema::volume::dsl; - let conn = - db_datastore.pool_connection_unauthorized().await.unwrap(); + let conn = datastore.pool_connection_unauthorized().await.unwrap(); let resources_to_clean_up = r#"{ "V1": { @@ -2867,14 +2865,14 @@ mod tests { // Soft delete the volume - let cr = db_datastore.soft_delete_volume(volume_id).await.unwrap(); + let cr = datastore.soft_delete_volume(volume_id).await.unwrap(); // Assert the contents of the returned CrucibleResources let datasets_and_regions = - db_datastore.regions_to_delete(&cr).await.unwrap(); + datastore.regions_to_delete(&cr).await.unwrap(); let datasets_and_snapshots = - db_datastore.snapshots_to_delete(&cr).await.unwrap(); + datastore.snapshots_to_delete(&cr).await.unwrap(); assert!(datasets_and_regions.is_empty()); assert_eq!(datasets_and_snapshots.len(), 1); @@ -2887,8 +2885,7 @@ mod tests { ); assert_eq!(region_snapshot.deleting, false); - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + datastore.terminate().await; logctx.cleanup_successful(); } @@ -2896,8 +2893,8 @@ mod tests { async fn test_volume_replace_region() { let logctx = dev::test_setup_log("test_volume_replace_region"); let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let (_opctx, db_datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&log).await; + let datastore = db.datastore(); // Insert four Region records (three, plus one additionally allocated) @@ -2912,7 +2909,7 @@ mod tests { ]; { - let conn = db_datastore.pool_connection_for_tests().await.unwrap(); + let conn = datastore.pool_connection_for_tests().await.unwrap(); for i in 0..4 { let (_, volume_id) = region_and_volume_ids[i]; @@ -2938,7 +2935,7 @@ mod tests { } } - let _volume = db_datastore + let _volume = datastore .volume_create(nexus_db_model::Volume::new( volume_id, serde_json::to_string(&VolumeConstructionRequest::Volume { @@ -2978,7 +2975,7 @@ mod tests { let target = region_and_volume_ids[0]; let replacement = region_and_volume_ids[3]; - let volume_replace_region_result = db_datastore + let volume_replace_region_result = datastore .volume_replace_region( /* target */ db::datastore::VolumeReplacementParams { @@ -3003,7 +3000,7 @@ mod tests { assert_eq!(volume_replace_region_result, VolumeReplaceResult::Done); let vcr: VolumeConstructionRequest = serde_json::from_str( - db_datastore.volume_get(volume_id).await.unwrap().unwrap().data(), + datastore.volume_get(volume_id).await.unwrap().unwrap().data(), ) .unwrap(); @@ -3040,7 +3037,7 @@ mod tests { ); // Now undo the replacement. Note volume ID is not swapped. - let volume_replace_region_result = db_datastore + let volume_replace_region_result = datastore .volume_replace_region( /* target */ db::datastore::VolumeReplacementParams { @@ -3065,7 +3062,7 @@ mod tests { assert_eq!(volume_replace_region_result, VolumeReplaceResult::Done); let vcr: VolumeConstructionRequest = serde_json::from_str( - db_datastore.volume_get(volume_id).await.unwrap().unwrap().data(), + datastore.volume_get(volume_id).await.unwrap().unwrap().data(), ) .unwrap(); @@ -3101,8 +3098,7 @@ mod tests { }, ); - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3110,8 +3106,8 @@ mod tests { async fn test_volume_replace_snapshot() { let logctx = dev::test_setup_log("test_volume_replace_snapshot"); let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let (_opctx, db_datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&log).await; + let datastore = db.datastore(); // Insert two volumes: one with the target to replace, and one temporary // "volume to delete" that's blank. @@ -3120,7 +3116,7 @@ mod tests { let volume_to_delete_id = Uuid::new_v4(); let rop_id = Uuid::new_v4(); - db_datastore + datastore .volume_create(nexus_db_model::Volume::new( volume_id, serde_json::to_string(&VolumeConstructionRequest::Volume { @@ -3179,7 +3175,7 @@ mod tests { .await .unwrap(); - db_datastore + datastore .volume_create(nexus_db_model::Volume::new( volume_to_delete_id, serde_json::to_string(&VolumeConstructionRequest::Volume { @@ -3195,7 +3191,7 @@ mod tests { // Do the replacement - let volume_replace_snapshot_result = db_datastore + let volume_replace_snapshot_result = datastore .volume_replace_snapshot( VolumeWithTarget(volume_id), ExistingTarget("[fd00:1122:3344:104::1]:400".parse().unwrap()), @@ -3212,7 +3208,7 @@ mod tests { // Ensure the shape of the resulting VCRs let vcr: VolumeConstructionRequest = serde_json::from_str( - db_datastore.volume_get(volume_id).await.unwrap().unwrap().data(), + datastore.volume_get(volume_id).await.unwrap().unwrap().data(), ) .unwrap(); @@ -3272,7 +3268,7 @@ mod tests { ); let vcr: VolumeConstructionRequest = serde_json::from_str( - db_datastore + datastore .volume_get(volume_to_delete_id) .await .unwrap() @@ -3313,7 +3309,7 @@ mod tests { // Now undo the replacement. Note volume ID is not swapped. - let volume_replace_snapshot_result = db_datastore + let volume_replace_snapshot_result = datastore .volume_replace_snapshot( VolumeWithTarget(volume_id), ExistingTarget("[fd55:1122:3344:101::1]:111".parse().unwrap()), @@ -3328,7 +3324,7 @@ mod tests { assert_eq!(volume_replace_snapshot_result, VolumeReplaceResult::Done,); let vcr: VolumeConstructionRequest = serde_json::from_str( - db_datastore.volume_get(volume_id).await.unwrap().unwrap().data(), + datastore.volume_get(volume_id).await.unwrap().unwrap().data(), ) .unwrap(); @@ -3389,7 +3385,7 @@ mod tests { ); let vcr: VolumeConstructionRequest = serde_json::from_str( - db_datastore + datastore .volume_get(volume_to_delete_id) .await .unwrap() @@ -3428,8 +3424,7 @@ mod tests { }, ); - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3438,14 +3433,14 @@ mod tests { let logctx = dev::test_setup_log("test_find_volumes_referencing_socket_addr"); let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let (opctx, db_datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let volume_id = Uuid::new_v4(); // case where the needle is found - db_datastore + datastore .volume_create(nexus_db_model::Volume::new( volume_id, serde_json::to_string(&VolumeConstructionRequest::Volume { @@ -3482,7 +3477,7 @@ mod tests { .await .unwrap(); - let volumes = db_datastore + let volumes = datastore .find_volumes_referencing_socket_addr( &opctx, "[fd00:1122:3344:104::1]:400".parse().unwrap(), @@ -3495,7 +3490,7 @@ mod tests { // case where the needle is missing - let volumes = db_datastore + let volumes = datastore .find_volumes_referencing_socket_addr( &opctx, "[fd55:1122:3344:104::1]:400".parse().unwrap(), @@ -3505,8 +3500,7 @@ mod tests { assert!(volumes.is_empty()); - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } diff --git a/nexus/db-queries/src/db/datastore/volume_repair.rs b/nexus/db-queries/src/db/datastore/volume_repair.rs index ba887048b4..115244f347 100644 --- a/nexus/db-queries/src/db/datastore/volume_repair.rs +++ b/nexus/db-queries/src/db/datastore/volume_repair.rs @@ -100,15 +100,14 @@ impl DataStore { mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use omicron_test_utils::dev; #[tokio::test] async fn volume_lock_conflict_error_returned() { let logctx = dev::test_setup_log("volume_lock_conflict_error_returned"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let lock_1 = Uuid::new_v4(); let lock_2 = Uuid::new_v4(); @@ -123,8 +122,7 @@ mod test { assert!(matches!(err, Error::Conflict { .. })); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/datastore/vpc.rs b/nexus/db-queries/src/db/datastore/vpc.rs index 2763aadf39..e3bd33e0a4 100644 --- a/nexus/db-queries/src/db/datastore/vpc.rs +++ b/nexus/db-queries/src/db/datastore/vpc.rs @@ -2755,9 +2755,9 @@ impl DataStore { #[cfg(test)] mod tests { use super::*; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::test::sled_baseboard_for_test; use crate::db::datastore::test::sled_system_hardware_for_test; - use crate::db::datastore::test_utils::datastore_test; use crate::db::datastore::test_utils::IneligibleSleds; use crate::db::model::Project; use crate::db::queries::vpc::MAX_VNI_SEARCH_RANGE_SIZE; @@ -2768,7 +2768,6 @@ mod tests { use nexus_reconfigurator_planning::blueprint_builder::BlueprintBuilder; use nexus_reconfigurator_planning::system::SledBuilder; use nexus_reconfigurator_planning::system::SystemDescription; - use nexus_test_utils::db::test_setup_database; use nexus_types::deployment::Blueprint; use nexus_types::deployment::BlueprintTarget; use nexus_types::deployment::BlueprintZoneConfig; @@ -2798,8 +2797,8 @@ mod tests { "test_project_create_vpc_raw_returns_none_on_vni_exhaustion", ); let log = &logctx.log; - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a project. let project_params = params::ProjectCreate { @@ -2889,8 +2888,7 @@ mod tests { else { panic!("Expected Ok(None) when creating a VPC without any available VNIs"); }; - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -2903,8 +2901,8 @@ mod tests { usdt::register_probes().unwrap(); let logctx = dev::test_setup_log("test_project_create_vpc_retries"); let log = &logctx.log; - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Create a project. let project_params = params::ProjectCreate { @@ -3000,8 +2998,7 @@ mod tests { } Err(e) => panic!("Unexpected error when inserting VPC: {e}"), }; - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3049,8 +3046,8 @@ mod tests { let logctx = dev::test_setup_log( "test_vpc_resolve_to_sleds_uses_current_target_blueprint", ); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Set up our fake system with 5 sleds. let rack_id = Uuid::new_v4(); @@ -3286,8 +3283,7 @@ mod tests { ) .await; - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3412,8 +3408,8 @@ mod tests { let logctx = dev::test_setup_log("test_vpc_system_router_sync_to_subnets"); let log = &logctx.log; - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (_, authz_vpc, db_vpc, _, db_router) = create_initial_vpc(log, &opctx, &datastore).await; @@ -3547,8 +3543,7 @@ mod tests { ) .await; - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -3640,8 +3635,8 @@ mod tests { let logctx = dev::test_setup_log("test_vpc_router_rule_instance_resolve"); let log = &logctx.log; - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let (authz_project, authz_vpc, db_vpc, authz_router, _) = create_initial_vpc(log, &opctx, &datastore).await; @@ -3777,8 +3772,7 @@ mod tests { _ => false, })); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/explain.rs b/nexus/db-queries/src/db/explain.rs index 823f064c91..284e96bc6e 100644 --- a/nexus/db-queries/src/db/explain.rs +++ b/nexus/db-queries/src/db/explain.rs @@ -94,10 +94,10 @@ mod test { use super::*; use crate::db; + use crate::db::datastore::pub_test_utils::TestDatabase; use async_bb8_diesel::AsyncSimpleConnection; use diesel::SelectableHelper; use expectorate::assert_contents; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; use uuid::Uuid; @@ -142,9 +142,8 @@ mod test { #[tokio::test] async fn test_explain_async() { let logctx = dev::test_setup_log("test_explain_async"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); create_schema(&pool).await; @@ -158,8 +157,7 @@ mod test { .unwrap(); assert_contents("tests/output/test-explain-output", &explanation); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -168,9 +166,8 @@ mod test { #[tokio::test] async fn test_explain_full_table_scan() { let logctx = dev::test_setup_log("test_explain_full_table_scan"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); create_schema(&pool).await; @@ -188,8 +185,7 @@ mod test { "Expected [{}] to contain 'FULL SCAN'", explanation ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/lookup.rs b/nexus/db-queries/src/db/lookup.rs index 6de950e935..43cd2a073f 100644 --- a/nexus/db-queries/src/db/lookup.rs +++ b/nexus/db-queries/src/db/lookup.rs @@ -958,24 +958,16 @@ mod test { use super::Instance; use super::LookupPath; use super::Project; - use crate::context::OpContext; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::model::Name; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; - use std::sync::Arc; /* This is a smoke test that things basically appear to work. */ #[tokio::test] async fn test_lookup() { let logctx = dev::test_setup_log("test_lookup"); - let mut db = test_setup_database(&logctx.log).await; - let (_, datastore) = - crate::db::datastore::test_utils::datastore_test(&logctx, &db) - .await; - let opctx = OpContext::for_tests( - logctx.log.new(o!()), - Arc::clone(&datastore) as Arc, - ); + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let project_name: Name = Name("my-project".parse().unwrap()); let instance_name: Name = Name("my-instance".parse().unwrap()); @@ -999,8 +991,7 @@ mod test { Project::PrimaryKey(_, p) if *p == project_id)); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/pagination.rs b/nexus/db-queries/src/db/pagination.rs index 785ac2ff54..a16591ad6c 100644 --- a/nexus/db-queries/src/db/pagination.rs +++ b/nexus/db-queries/src/db/pagination.rs @@ -343,11 +343,11 @@ mod test { use super::*; use crate::db; + use crate::db::datastore::pub_test_utils::TestDatabase; use async_bb8_diesel::{AsyncRunQueryDsl, AsyncSimpleConnection}; use diesel::JoinOnDsl; use diesel::SelectableHelper; use dropshot::PaginationOrder; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external::DataPageParams; use omicron_test_utils::dev; use std::num::NonZeroU32; @@ -489,9 +489,8 @@ mod test { async fn test_paginated_single_column_ascending() { let logctx = dev::test_setup_log("test_paginated_single_column_ascending"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); use schema::test_users::dsl; @@ -516,8 +515,7 @@ mod test { let observed = execute_query(&pool, query).await; assert_eq!(observed, vec![(2, 2), (3, 3)]); - pool.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -525,9 +523,8 @@ mod test { async fn test_paginated_single_column_descending() { let logctx = dev::test_setup_log("test_paginated_single_column_descending"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); use schema::test_users::dsl; @@ -552,8 +549,7 @@ mod test { let observed = execute_query(&pool, query).await; assert_eq!(observed, vec![(2, 2), (1, 1)]); - pool.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -561,9 +557,8 @@ mod test { async fn test_paginated_multicolumn_ascending() { let logctx = dev::test_setup_log("test_paginated_multicolumn_ascending"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); use schema::test_users::dsl; @@ -607,8 +602,7 @@ mod test { let observed = execute_query(&pool, query).await; assert_eq!(observed, vec![(1, 1), (2, 1), (3, 1), (1, 2), (2, 3)]); - pool.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -616,9 +610,8 @@ mod test { async fn test_paginated_multicolumn_descending() { let logctx = dev::test_setup_log("test_paginated_multicolumn_descending"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); use schema::test_users::dsl; @@ -662,8 +655,7 @@ mod test { let observed = execute_query(&pool, query).await; assert_eq!(observed, vec![(2, 3), (1, 2), (3, 1), (2, 1), (1, 1)]); - pool.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } @@ -673,9 +665,8 @@ mod test { let logctx = dev::test_setup_log("test_paginated_multicolumn_works_with_joins"); - let mut db = test_setup_database(&logctx.log).await; - let cfg = db::Config { url: db.pg_config().clone() }; - let pool = db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); use schema::test_phone_numbers::dsl as phone_numbers_dsl; use schema::test_users::dsl; @@ -764,8 +755,7 @@ mod test { &[((2, 3), 42), ((3, 1), 50), ((3, 1), 51), ((3, 1), 52)] ); - pool.terminate().await; - let _ = db.cleanup().await; + db.terminate().await; logctx.cleanup_successful(); } diff --git a/nexus/db-queries/src/db/queries/external_ip.rs b/nexus/db-queries/src/db/queries/external_ip.rs index bd0b5be5ac..2f17a5e2b6 100644 --- a/nexus/db-queries/src/db/queries/external_ip.rs +++ b/nexus/db-queries/src/db/queries/external_ip.rs @@ -865,8 +865,7 @@ impl RunQueryDsl for NextExternalIp {} #[cfg(test)] mod tests { use crate::authz; - use crate::context::OpContext; - use crate::db::datastore::DataStore; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::SERVICE_IP_POOL_NAME; use crate::db::identity::Resource; use crate::db::lookup::LookupPath; @@ -882,7 +881,6 @@ mod tests { use nexus_db_model::IpPoolResource; use nexus_db_model::IpPoolResourceType; use nexus_sled_agent_shared::inventory::ZoneKind; - use nexus_test_utils::db::test_setup_database; use nexus_types::deployment::OmicronZoneExternalFloatingIp; use nexus_types::deployment::OmicronZoneExternalIp; use nexus_types::deployment::OmicronZoneExternalSnatIp; @@ -893,32 +891,25 @@ mod tests { use omicron_common::api::external::Error; use omicron_common::api::external::IdentityMetadataCreateParams; use omicron_test_utils::dev; - use omicron_test_utils::dev::db::CockroachInstance; use omicron_uuid_kinds::ExternalIpUuid; use omicron_uuid_kinds::GenericUuid; use omicron_uuid_kinds::InstanceUuid; use omicron_uuid_kinds::OmicronZoneUuid; use std::net::IpAddr; use std::net::Ipv4Addr; - use std::sync::Arc; use uuid::Uuid; struct TestContext { logctx: LogContext, - opctx: OpContext, - db: CockroachInstance, - db_datastore: Arc, + db: TestDatabase, } impl TestContext { async fn new(test_name: &str) -> Self { let logctx = dev::test_setup_log(test_name); let log = logctx.log.new(o!()); - let db = test_setup_database(&log).await; - let (opctx, db_datastore) = - crate::db::datastore::test_utils::datastore_test(&logctx, &db) - .await; - Self { logctx, opctx, db, db_datastore } + let db = TestDatabase::new_with_datastore(&log).await; + Self { logctx, db } } /// Create pool, associate with current silo @@ -933,26 +924,28 @@ mod tests { description: format!("ip pool {}", name), }); - self.db_datastore - .ip_pool_create(&self.opctx, pool.clone()) + self.db + .datastore() + .ip_pool_create(self.db.opctx(), pool.clone()) .await .expect("Failed to create IP pool"); - let silo_id = self.opctx.authn.silo_required().unwrap().id(); + let silo_id = self.db.opctx().authn.silo_required().unwrap().id(); let association = IpPoolResource { resource_id: silo_id, resource_type: IpPoolResourceType::Silo, ip_pool_id: pool.id(), is_default, }; - self.db_datastore - .ip_pool_link_silo(&self.opctx, association) + self.db + .datastore() + .ip_pool_link_silo(self.db.opctx(), association) .await .expect("Failed to associate IP pool with silo"); self.initialize_ip_pool(name, range).await; - LookupPath::new(&self.opctx, &self.db_datastore) + LookupPath::new(self.db.opctx(), &self.db.datastore()) .ip_pool_id(pool.id()) .lookup_for(authz::Action::Read) .await @@ -964,8 +957,9 @@ mod tests { // Find the target IP pool use crate::db::schema::ip_pool::dsl as ip_pool_dsl; let conn = self - .db_datastore - .pool_connection_authorized(&self.opctx) + .db + .datastore() + .pool_connection_authorized(self.db.opctx()) .await .unwrap(); let pool = ip_pool_dsl::ip_pool @@ -984,8 +978,9 @@ mod tests { .values(pool_range) .execute_async( &*self - .db_datastore - .pool_connection_authorized(&self.opctx) + .db + .datastore() + .pool_connection_authorized(self.db.opctx()) .await .unwrap(), ) @@ -1012,8 +1007,9 @@ mod tests { }); let conn = self - .db_datastore - .pool_connection_authorized(&self.opctx) + .db + .datastore() + .pool_connection_authorized(self.db.opctx()) .await .unwrap(); @@ -1029,16 +1025,16 @@ mod tests { async fn default_pool_id(&self) -> Uuid { let (.., pool) = self - .db_datastore - .ip_pools_fetch_default(&self.opctx) + .db + .datastore() + .ip_pools_fetch_default(self.db.opctx()) .await .expect("Failed to lookup default ip pool"); pool.identity.id } - async fn success(mut self) { - self.db_datastore.terminate().await; - self.db.cleanup().await.unwrap(); + async fn success(self) { + self.db.terminate().await; self.logctx.cleanup_successful(); } } @@ -1060,9 +1056,10 @@ mod tests { let id = Uuid::new_v4(); let instance_id = InstanceUuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), id, instance_id, context.default_pool_id().await, @@ -1077,9 +1074,10 @@ mod tests { // The next allocation should fail, due to IP exhaustion let instance_id = InstanceUuid::new_v4(); let err = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, context.default_pool_id().await, @@ -1114,9 +1112,10 @@ mod tests { // the only address in the pool. let instance_id = context.create_instance("for-eph").await; let ephemeral_ip = context - .db_datastore + .db + .datastore() .allocate_instance_ephemeral_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, /* pool_name = */ None, @@ -1133,9 +1132,10 @@ mod tests { // nor any SNAT IPs. let instance_id = context.create_instance("for-snat").await; let res = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, context.default_pool_id().await, @@ -1156,9 +1156,10 @@ mod tests { ); let res = context - .db_datastore + .db + .datastore() .allocate_instance_ephemeral_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, /* pool_name = */ None, @@ -1210,9 +1211,10 @@ mod tests { for (expected_ip, expected_first_port) in external_ips.clone().take(2) { let instance_id = InstanceUuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, context.default_pool_id().await, @@ -1229,8 +1231,9 @@ mod tests { // Release the first context - .db_datastore - .deallocate_external_ip(&context.opctx, ips[0].id) + .db + .datastore() + .deallocate_external_ip(&context.db.opctx(), ips[0].id) .await .expect("Failed to release the first external IP address"); @@ -1238,9 +1241,10 @@ mod tests { // released. let instance_id = InstanceUuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, context.default_pool_id().await, @@ -1265,9 +1269,10 @@ mod tests { // from the original loop. let instance_id = InstanceUuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, context.default_pool_id().await, @@ -1302,9 +1307,10 @@ mod tests { let pool_name = None; let ip = context - .db_datastore + .db + .datastore() .allocate_instance_ephemeral_ip( - &context.opctx, + &context.db.opctx(), id, instance_id, pool_name, @@ -1350,9 +1356,10 @@ mod tests { // service. let service_id = OmicronZoneUuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), service_id, ZoneKind::Nexus, ip_10_0_0_3, @@ -1367,9 +1374,10 @@ mod tests { // Try allocating the same service IP again. let ip_again = context - .db_datastore + .db + .datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), service_id, ZoneKind::Nexus, ip_10_0_0_3, @@ -1383,9 +1391,9 @@ mod tests { // Try allocating the same service IP once more, but do it with a // different UUID. let err = context - .db_datastore + .db.datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), service_id, ZoneKind::Nexus, OmicronZoneExternalIp::Floating(OmicronZoneExternalFloatingIp { @@ -1403,9 +1411,9 @@ mod tests { // Try allocating the same service IP once more, but do it with a // different input address. let err = context - .db_datastore + .db.datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), service_id, ZoneKind::Nexus, OmicronZoneExternalIp::Floating(OmicronZoneExternalFloatingIp { @@ -1429,9 +1437,9 @@ mod tests { .unwrap(), }); let err = context - .db_datastore + .db.datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), service_id, ZoneKind::BoundaryNtp, ip_10_0_0_3_snat_0, @@ -1456,9 +1464,10 @@ mod tests { }); let snat_service_id = OmicronZoneUuid::new_v4(); let snat_ip = context - .db_datastore + .db + .datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), snat_service_id, ZoneKind::BoundaryNtp, ip_10_0_0_1_snat_32768, @@ -1477,9 +1486,10 @@ mod tests { // Try allocating the same service IP again. let snat_ip_again = context - .db_datastore + .db + .datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), snat_service_id, ZoneKind::BoundaryNtp, ip_10_0_0_1_snat_32768, @@ -1505,9 +1515,9 @@ mod tests { .unwrap(), }); let err = context - .db_datastore + .db.datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), snat_service_id, ZoneKind::BoundaryNtp, ip_10_0_0_1_snat_49152, @@ -1544,9 +1554,10 @@ mod tests { let service_id = OmicronZoneUuid::new_v4(); let err = context - .db_datastore + .db + .datastore() .external_ip_allocate_omicron_zone( - &context.opctx, + &context.db.opctx(), service_id, ZoneKind::Nexus, ip_10_0_0_5, @@ -1578,9 +1589,10 @@ mod tests { let instance_id = InstanceUuid::new_v4(); let id = Uuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), id, instance_id, context.default_pool_id().await, @@ -1595,9 +1607,10 @@ mod tests { // Create a new IP, with the _same_ ID, and ensure we get back the same // value. let new_ip = context - .db_datastore + .db + .datastore() .allocate_instance_snat_ip( - &context.opctx, + &context.db.opctx(), id, instance_id, context.default_pool_id().await, @@ -1648,9 +1661,10 @@ mod tests { let id = Uuid::new_v4(); let ip = context - .db_datastore + .db + .datastore() .allocate_instance_ephemeral_ip( - &context.opctx, + &context.db.opctx(), id, instance_id, Some(p1), @@ -1693,9 +1707,10 @@ mod tests { let instance_id = context.create_instance(&format!("o{octet}")).await; let ip = context - .db_datastore + .db + .datastore() .allocate_instance_ephemeral_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, Some(p1.clone()), @@ -1715,9 +1730,10 @@ mod tests { // Allocating another address should _fail_, and not use the first pool. let instance_id = context.create_instance("final").await; context - .db_datastore + .db + .datastore() .allocate_instance_ephemeral_ip( - &context.opctx, + &context.db.opctx(), Uuid::new_v4(), instance_id, Some(p1), diff --git a/nexus/db-queries/src/db/queries/network_interface.rs b/nexus/db-queries/src/db/queries/network_interface.rs index 01aa8aafe2..39c799d223 100644 --- a/nexus/db-queries/src/db/queries/network_interface.rs +++ b/nexus/db-queries/src/db/queries/network_interface.rs @@ -1797,6 +1797,7 @@ mod tests { use super::NUM_INITIAL_RESERVED_IP_ADDRESSES; use crate::authz; use crate::context::OpContext; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::DataStore; use crate::db::identity::Resource; use crate::db::lookup::LookupPath; @@ -1811,7 +1812,6 @@ mod tests { use async_bb8_diesel::AsyncRunQueryDsl; use dropshot::test_util::LogContext; use model::NetworkInterfaceKind; - use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::params; use nexus_types::external_api::params::InstanceCreate; use nexus_types::external_api::params::InstanceNetworkInterfaceAttachment; @@ -1822,7 +1822,6 @@ mod tests { use omicron_common::api::external::InstanceCpuCount; use omicron_common::api::external::MacAddr; use omicron_test_utils::dev; - use omicron_test_utils::dev::db::CockroachInstance; use omicron_uuid_kinds::GenericUuid; use omicron_uuid_kinds::InstanceUuid; use oxnet::Ipv4Net; @@ -1832,7 +1831,6 @@ mod tests { use std::net::IpAddr; use std::net::Ipv4Addr; use std::net::Ipv6Addr; - use std::sync::Arc; use uuid::Uuid; // Add an instance. We'll use this to verify that the instance must be @@ -1964,9 +1962,7 @@ mod tests { // Context for testing network interface queries. struct TestContext { logctx: LogContext, - opctx: OpContext, - db: CockroachInstance, - db_datastore: Arc, + db: TestDatabase, project_id: Uuid, net1: Network, net2: Network, @@ -1976,10 +1972,8 @@ mod tests { async fn new(test_name: &str, n_subnets: u8) -> Self { let logctx = dev::test_setup_log(test_name); let log = logctx.log.new(o!()); - let db = test_setup_database(&log).await; - let (opctx, db_datastore) = - crate::db::datastore::test_utils::datastore_test(&logctx, &db) - .await; + let db = TestDatabase::new_with_datastore(&log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); let authz_silo = opctx.authn.silo_required().unwrap(); @@ -1994,11 +1988,11 @@ mod tests { }, ); let (.., project) = - db_datastore.project_create(&opctx, project).await.unwrap(); + datastore.project_create(&opctx, project).await.unwrap(); use crate::db::schema::vpc_subnet::dsl::vpc_subnet; let conn = - db_datastore.pool_connection_authorized(&opctx).await.unwrap(); + datastore.pool_connection_authorized(&opctx).await.unwrap(); let net1 = Network::new(n_subnets); let net2 = Network::new(n_subnets); for subnet in net1.subnets.iter().chain(net2.subnets.iter()) { @@ -2009,30 +2003,29 @@ mod tests { .unwrap(); } drop(conn); - Self { - logctx, - opctx, - db, - db_datastore, - project_id: project.id(), - net1, - net2, - } + Self { logctx, db, project_id: project.id(), net1, net2 } + } + + fn opctx(&self) -> &OpContext { + self.db.opctx() + } + + fn datastore(&self) -> &DataStore { + self.db.datastore() } - async fn success(mut self) { - self.db_datastore.terminate().await; - self.db.cleanup().await.unwrap(); + async fn success(self) { + self.db.terminate().await; self.logctx.cleanup_successful(); } async fn create_stopped_instance(&self) -> Instance { instance_set_state( - &self.db_datastore, + self.datastore(), create_instance( - &self.opctx, + self.opctx(), self.project_id, - &self.db_datastore, + self.datastore(), ) .await, InstanceState::NoVmm, @@ -2042,11 +2035,11 @@ mod tests { async fn create_running_instance(&self) -> Instance { instance_set_state( - &self.db_datastore, + self.datastore(), create_instance( - &self.opctx, + self.opctx(), self.project_id, - &self.db_datastore, + self.datastore(), ) .await, InstanceState::Vmm, @@ -2079,17 +2072,17 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, interface) + .datastore() + .service_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); // We should be able to delete twice, and be told that the first delete // modified the row and the second did not. let first_deleted = context - .db_datastore + .datastore() .service_delete_network_interface( - &context.opctx, + context.opctx(), service_id, inserted_interface.id(), ) @@ -2098,9 +2091,9 @@ mod tests { assert!(first_deleted, "first delete removed interface"); let second_deleted = context - .db_datastore + .datastore() .service_delete_network_interface( - &context.opctx, + context.opctx(), service_id, inserted_interface.id(), ) @@ -2111,9 +2104,9 @@ mod tests { // Attempting to delete a nonexistent interface should fail. let bogus_id = Uuid::new_v4(); let err = context - .db_datastore + .datastore() .service_delete_network_interface( - &context.opctx, + context.opctx(), service_id, bogus_id, ) @@ -2148,8 +2141,8 @@ mod tests { Some(requested_ip), ) .unwrap(); - let err = context.db_datastore - .instance_create_network_interface_raw(&context.opctx, interface.clone()) + let err = context.datastore() + .instance_create_network_interface_raw(context.opctx(), interface.clone()) .await .expect_err("Should not be able to create an interface for a running instance"); assert!( @@ -2178,9 +2171,9 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface.clone(), ) .await @@ -2209,8 +2202,8 @@ mod tests { None, ) .unwrap(); - let err = context.db_datastore - .instance_create_network_interface_raw(&context.opctx, interface.clone()) + let err = context.datastore() + .instance_create_network_interface_raw(context.opctx(), interface.clone()) .await .expect_err("Should not be able to insert an interface for an instance that doesn't exist"); assert!( @@ -2248,9 +2241,9 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface.clone(), ) .await @@ -2292,8 +2285,8 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); @@ -2311,8 +2304,8 @@ mod tests { ) .unwrap(); let result = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await; assert!( matches!(result, Err(InsertError::IpAddressNotAvailable(_))), @@ -2348,8 +2341,8 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, interface) + .datastore() + .service_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); assert_eq!(inserted_interface.mac.0, mac); @@ -2388,8 +2381,11 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, interface) + .datastore() + .service_create_network_interface_raw( + context.opctx(), + interface, + ) .await .expect("Failed to insert interface"); assert_eq!(*inserted_interface.slot, slot); @@ -2425,8 +2421,8 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, interface) + .datastore() + .service_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); assert_eq!(inserted_interface.mac.0, mac); @@ -2448,8 +2444,11 @@ mod tests { ) .unwrap(); let result = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, new_interface) + .datastore() + .service_create_network_interface_raw( + context.opctx(), + new_interface, + ) .await; assert!( matches!(result, Err(InsertError::MacAddressNotAvailable(_))), @@ -2501,8 +2500,8 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, interface) + .datastore() + .service_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); assert_eq!(*inserted_interface.slot, 0); @@ -2522,8 +2521,11 @@ mod tests { ) .unwrap(); let result = context - .db_datastore - .service_create_network_interface_raw(&context.opctx, new_interface) + .datastore() + .service_create_network_interface_raw( + context.opctx(), + new_interface, + ) .await; assert!( matches!(result, Err(InsertError::SlotNotAvailable(0))), @@ -2550,9 +2552,9 @@ mod tests { ) .unwrap(); let _ = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface.clone(), ) .await @@ -2569,8 +2571,8 @@ mod tests { ) .unwrap(); let result = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await; assert!( matches!( @@ -2600,8 +2602,8 @@ mod tests { ) .unwrap(); let _ = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); let interface = IncompleteNetworkInterface::new_instance( @@ -2616,8 +2618,8 @@ mod tests { ) .unwrap(); let result = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await; assert!( matches!(result, Err(InsertError::NonUniqueVpcSubnets)), @@ -2644,16 +2646,16 @@ mod tests { ) .unwrap(); let _ = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface.clone(), ) .await .expect("Failed to insert interface"); let result = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await; assert!( matches!( @@ -2686,8 +2688,8 @@ mod tests { ) .unwrap(); let _ = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await .expect("Failed to insert interface"); let expected_address = "172.30.0.5".parse().unwrap(); @@ -2704,9 +2706,9 @@ mod tests { ) .unwrap(); let result = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface, ) .await; @@ -2740,9 +2742,9 @@ mod tests { ) .unwrap(); let _ = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface, ) .await @@ -2751,9 +2753,9 @@ mod tests { // Next one should fail let instance = create_stopped_instance( - &context.opctx, + context.opctx(), context.project_id, - &context.db_datastore, + context.datastore(), ) .await; let instance_id = InstanceUuid::from_untyped_uuid(instance.id()); @@ -2769,8 +2771,8 @@ mod tests { ) .unwrap(); let result = context - .db_datastore - .instance_create_network_interface_raw(&context.opctx, interface) + .datastore() + .instance_create_network_interface_raw(context.opctx(), interface) .await; assert!( matches!(result, Err(InsertError::NoAvailableIpAddresses)), @@ -2802,9 +2804,9 @@ mod tests { ) .unwrap(); let result = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface, ) .await; @@ -2870,9 +2872,9 @@ mod tests { ) .unwrap(); let inserted_interface = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface.clone(), ) .await @@ -2905,9 +2907,9 @@ mod tests { ) .unwrap(); let result = context - .db_datastore + .datastore() .instance_create_network_interface_raw( - &context.opctx, + context.opctx(), interface.clone(), ) .await diff --git a/nexus/db-queries/src/db/queries/next_item.rs b/nexus/db-queries/src/db/queries/next_item.rs index 3a65546e7c..0ec0727737 100644 --- a/nexus/db-queries/src/db/queries/next_item.rs +++ b/nexus/db-queries/src/db/queries/next_item.rs @@ -924,6 +924,7 @@ mod tests { use super::NextItem; use super::ShiftIndices; use crate::db; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::explain::ExplainableAsync as _; use crate::db::queries::next_item::NextItemSelfJoined; use async_bb8_diesel::AsyncRunQueryDsl; @@ -937,9 +938,7 @@ mod tests { use diesel::Column; use diesel::Insertable; use diesel::SelectableHelper; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; - use std::sync::Arc; use uuid::Uuid; table! { @@ -1102,11 +1101,8 @@ mod tests { async fn test_wrapping_next_item_query() { // Setup the test database let logctx = dev::test_setup_log("test_wrapping_next_item_query"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // We're going to operate on a separate table, for simplicity. @@ -1156,8 +1152,7 @@ mod tests { .unwrap(); assert_eq!(it.value, 2); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1166,11 +1161,8 @@ mod tests { // Setup the test database let logctx = dev::test_setup_log("test_next_item_query_is_ordered_by_indices"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // We're going to operate on a separate table, for simplicity. @@ -1213,8 +1205,7 @@ mod tests { "The next item query should not have further items to generate", ); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1257,11 +1248,8 @@ mod tests { async fn test_explain_next_item_self_joined() { // Setup the test database let logctx = dev::test_setup_log("test_explain_next_item_self_joined"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // We're going to operate on a separate table, for simplicity. @@ -1276,8 +1264,7 @@ mod tests { >::new_scoped(Uuid::nil(), i32::MIN, i32::MAX); let out = query.explain_async(&conn).await.unwrap(); println!("{out}"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1285,11 +1272,8 @@ mod tests { async fn test_next_item_self_joined() { // Setup the test database let logctx = dev::test_setup_log("test_next_item_self_joined"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // We're going to operate on a separate table, for simplicity. @@ -1313,8 +1297,7 @@ mod tests { .get_result_async(&*conn) .await .expect_err("should not be able to insert after the query range is exhausted"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1323,11 +1306,8 @@ mod tests { // Setup the test database let logctx = dev::test_setup_log("test_next_item_self_joined_with_gaps"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // We're going to operate on a separate table, for simplicity. @@ -1368,8 +1348,7 @@ mod tests { "Should have inserted the next skipped value" ); } - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -1378,11 +1357,8 @@ mod tests { async fn print_next_item_query_forms() { // Setup the test database let logctx = dev::test_setup_log("print_next_item_query_forms"); - let log = logctx.log.new(o!()); - let db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); // We're going to operate on a separate table, for simplicity. diff --git a/nexus/db-queries/src/db/queries/oximeter.rs b/nexus/db-queries/src/db/queries/oximeter.rs index 4a4bf97235..eee7dd5669 100644 --- a/nexus/db-queries/src/db/queries/oximeter.rs +++ b/nexus/db-queries/src/db/queries/oximeter.rs @@ -221,9 +221,9 @@ pub fn reassign_producers_query(oximeter_id: Uuid) -> TypedSqlQuery<()> { #[cfg(test)] mod test { use super::*; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::explain::ExplainableAsync; use crate::db::raw_query_builder::expectorate_query_contents; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; use std::time::Duration; use uuid::Uuid; @@ -266,10 +266,8 @@ mod test { #[tokio::test] async fn explainable_upsert_producer() { let logctx = dev::test_setup_log("explainable_upsert_producer"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let producer = internal::nexus::ProducerEndpoint { @@ -285,18 +283,15 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn explainable_reassign_producers() { let logctx = dev::test_setup_log("explainable_reassign_producers"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let oximeter_id = Uuid::nil(); @@ -307,8 +302,7 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/queries/region_allocation.rs b/nexus/db-queries/src/db/queries/region_allocation.rs index c7e23be9db..a531e726b3 100644 --- a/nexus/db-queries/src/db/queries/region_allocation.rs +++ b/nexus/db-queries/src/db/queries/region_allocation.rs @@ -405,10 +405,10 @@ UNION #[cfg(test)] mod test { use super::*; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::datastore::REGION_REDUNDANCY_THRESHOLD; use crate::db::explain::ExplainableAsync; use crate::db::raw_query_builder::expectorate_query_contents; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; use uuid::Uuid; @@ -504,10 +504,8 @@ mod test { #[tokio::test] async fn explainable() { let logctx = dev::test_setup_log("explainable"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let volume_id = Uuid::new_v4(); @@ -546,8 +544,7 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs index 2784afe9c8..8d3ef320ee 100644 --- a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs +++ b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs @@ -478,9 +478,9 @@ FROM #[cfg(test)] mod test { use super::*; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::explain::ExplainableAsync; use crate::db::raw_query_builder::expectorate_query_contents; - use nexus_test_utils::db::test_setup_database; use omicron_test_utils::dev; use uuid::Uuid; @@ -565,10 +565,8 @@ mod test { #[tokio::test] async fn explain_insert_storage() { let logctx = dev::test_setup_log("explain_insert_storage"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let id = Uuid::nil(); @@ -587,18 +585,15 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn explain_delete_storage() { let logctx = dev::test_setup_log("explain_delete_storage"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let id = Uuid::nil(); @@ -615,18 +610,15 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn explain_insert_instance() { let logctx = dev::test_setup_log("explain_insert_instance"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let id = InstanceUuid::nil(); @@ -642,18 +634,15 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } #[tokio::test] async fn explain_delete_instance() { let logctx = dev::test_setup_log("explain_delete_instance"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = crate::db::Pool::new_single_host(&logctx.log, &cfg); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); let conn = pool.claim().await.unwrap(); let id = InstanceUuid::nil(); @@ -669,8 +658,7 @@ mod test { .await .expect("Failed to explain query - is it valid SQL?"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/db/queries/vpc_subnet.rs b/nexus/db-queries/src/db/queries/vpc_subnet.rs index ff34c50e51..17a362880f 100644 --- a/nexus/db-queries/src/db/queries/vpc_subnet.rs +++ b/nexus/db-queries/src/db/queries/vpc_subnet.rs @@ -288,14 +288,13 @@ impl InsertVpcSubnetError { mod test { use super::InsertVpcSubnetError; use super::InsertVpcSubnetQuery; + use crate::db::datastore::pub_test_utils::TestDatabase; use crate::db::explain::ExplainableAsync as _; use crate::db::model::VpcSubnet; - use nexus_test_utils::db::test_setup_database; use omicron_common::api::external::IdentityMetadataCreateParams; use omicron_common::api::external::Name; use omicron_test_utils::dev; use std::convert::TryInto; - use std::sync::Arc; #[tokio::test] async fn explain_insert_query() { @@ -310,16 +309,11 @@ mod test { VpcSubnet::new(subnet_id, vpc_id, identity, ipv4_block, ipv6_block); let query = InsertVpcSubnetQuery::new(row); let logctx = dev::test_setup_log("explain_insert_query"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); - let conn = pool.claim().await.unwrap(); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let conn = db.pool().claim().await.unwrap(); let explain = query.explain_async(&conn).await.unwrap(); println!("{explain}"); - pool.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -351,16 +345,8 @@ mod test { // Setup the test database let logctx = dev::test_setup_log("test_insert_vpc_subnet_query"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); - let db_datastore = Arc::new( - crate::db::DataStore::new(&log, Arc::clone(&pool), None) - .await - .unwrap(), - ); + let db = TestDatabase::new_with_raw_datastore(&logctx.log).await; + let db_datastore = db.datastore(); // We should be able to insert anything into an empty table. assert!( @@ -471,8 +457,7 @@ mod test { "Should be able to insert new VPC Subnet with non-overlapping IP ranges" ); - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -545,16 +530,8 @@ mod test { // Setup the test database let logctx = dev::test_setup_log("test_insert_vpc_subnet_query_is_idempotent"); - let log = logctx.log.new(o!()); - let mut db = test_setup_database(&log).await; - let cfg = crate::db::Config { url: db.pg_config().clone() }; - let pool = - Arc::new(crate::db::Pool::new_single_host(&logctx.log, &cfg)); - let db_datastore = Arc::new( - crate::db::DataStore::new(&log, Arc::clone(&pool), None) - .await - .unwrap(), - ); + let db = TestDatabase::new_with_raw_datastore(&logctx.log).await; + let db_datastore = db.datastore(); // We should be able to insert anything into an empty table. let inserted = db_datastore @@ -573,8 +550,7 @@ mod test { "Must be able to insert the exact same VPC subnet more than once", ); assert_rows_eq(&inserted, &row); - db_datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/db-queries/src/policy_test/mod.rs b/nexus/db-queries/src/policy_test/mod.rs index a14ffb943b..be81f58138 100644 --- a/nexus/db-queries/src/policy_test/mod.rs +++ b/nexus/db-queries/src/policy_test/mod.rs @@ -14,7 +14,7 @@ mod coverage; mod resource_builder; mod resources; -use crate::db; +use crate::db::datastore::pub_test_utils::TestDatabase; use coverage::Coverage; use futures::StreamExt; use nexus_auth::authn; @@ -22,7 +22,6 @@ use nexus_auth::authn::SiloAuthnPolicy; use nexus_auth::authn::USER_TEST_PRIVILEGED; use nexus_auth::authz; use nexus_auth::context::OpContext; -use nexus_test_utils::db::test_setup_database; use nexus_types::external_api::shared; use nexus_types::external_api::shared::FleetRole; use nexus_types::external_api::shared::SiloRole; @@ -62,9 +61,8 @@ use uuid::Uuid; #[tokio::test(flavor = "multi_thread")] async fn test_iam_roles_behavior() { let logctx = dev::test_setup_log("test_iam_roles"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - db::datastore::test_utils::datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Before we can create the resources, users, and role assignments that we // need, we must grant the "test-privileged" user privileges to fetch and @@ -173,8 +171,7 @@ async fn test_iam_roles_behavior() { &std::str::from_utf8(buffer.as_ref()).expect("non-UTF8 output"), ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -330,9 +327,8 @@ impl Write for StdoutTee { async fn test_conferred_roles() { // To start, this test looks a lot like the test above. let logctx = dev::test_setup_log("test_conferred_roles"); - let mut db = test_setup_database(&logctx.log).await; - let (opctx, datastore) = - db::datastore::test_utils::datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let (opctx, datastore) = (db.opctx(), db.datastore()); // Before we can create the resources, users, and role assignments that we // need, we must grant the "test-privileged" user privileges to fetch and @@ -465,7 +461,6 @@ async fn test_conferred_roles() { &std::str::from_utf8(buffer.as_ref()).expect("non-UTF8 output"), ); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } diff --git a/nexus/db-queries/src/transaction_retry.rs b/nexus/db-queries/src/transaction_retry.rs index f8c59ee15b..9975d1cf3b 100644 --- a/nexus/db-queries/src/transaction_retry.rs +++ b/nexus/db-queries/src/transaction_retry.rs @@ -258,8 +258,7 @@ impl OptionalError { mod test { use super::*; - use crate::db::datastore::test_utils::datastore_test; - use nexus_test_utils::db::test_setup_database; + use crate::db::datastore::pub_test_utils::TestDatabase; use omicron_test_utils::dev; use oximeter::types::FieldValue; @@ -271,8 +270,8 @@ mod test { let logctx = dev::test_setup_log( "test_transaction_rollback_produces_no_samples", ); - let mut db = test_setup_database(&logctx.log).await; - let (_opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let conn = datastore.pool_connection_for_tests().await.unwrap(); @@ -294,8 +293,7 @@ mod test { .clone(); assert_eq!(samples, vec![]); - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } @@ -305,8 +303,8 @@ mod test { async fn test_transaction_retry_produces_samples() { let logctx = dev::test_setup_log("test_transaction_retry_produces_samples"); - let mut db = test_setup_database(&logctx.log).await; - let (_opctx, datastore) = datastore_test(&logctx, &db).await; + let db = TestDatabase::new_with_datastore(&logctx.log).await; + let datastore = db.datastore(); let conn = datastore.pool_connection_for_tests().await.unwrap(); datastore @@ -356,8 +354,7 @@ mod test { ); } - datastore.terminate().await; - db.cleanup().await.unwrap(); + db.terminate().await; logctx.cleanup_successful(); } } diff --git a/nexus/metrics-producer-gc/src/lib.rs b/nexus/metrics-producer-gc/src/lib.rs index 5785af834f..32e2be5809 100644 --- a/nexus/metrics-producer-gc/src/lib.rs +++ b/nexus/metrics-producer-gc/src/lib.rs @@ -218,7 +218,7 @@ mod tests { let logctx = dev::test_setup_log("test_prune_expired_producers"); let mut db = test_setup_database(&logctx.log).await; let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + datastore_test(&logctx.log, &db, Uuid::new_v4()).await; // Insert an Oximeter collector let collector_info = OximeterInfo::new(¶ms::OximeterInfo { @@ -304,7 +304,7 @@ mod tests { ); let mut db = test_setup_database(&logctx.log).await; let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + datastore_test(&logctx.log, &db, Uuid::new_v4()).await; let mut collector = httptest::Server::run(); diff --git a/nexus/src/app/background/tasks/crdb_node_id_collector.rs b/nexus/src/app/background/tasks/crdb_node_id_collector.rs index 0867aa9e17..ddb068226f 100644 --- a/nexus/src/app/background/tasks/crdb_node_id_collector.rs +++ b/nexus/src/app/background/tasks/crdb_node_id_collector.rs @@ -349,7 +349,7 @@ mod tests { let logctx = dev::test_setup_log("test_activate_fails_if_no_blueprint"); let mut db = test_setup_database(&logctx.log).await; let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + datastore_test(&logctx.log, &db, Uuid::new_v4()).await; let (_tx_blueprint, rx_blueprint) = watch::channel(None); let mut collector = @@ -381,7 +381,7 @@ mod tests { dev::test_setup_log("test_activate_with_no_unknown_node_ids"); let mut db = test_setup_database(&logctx.log).await; let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + datastore_test(&logctx.log, &db, Uuid::new_v4()).await; let blueprint = BlueprintBuilder::build_empty_with_sleds( iter::once(SledUuid::new_v4()), @@ -446,7 +446,7 @@ mod tests { let logctx = dev::test_setup_log("test_activate_with_unknown_node_ids"); let mut db = test_setup_database(&logctx.log).await; let (opctx, datastore) = - datastore_test(&logctx, &db, Uuid::new_v4()).await; + datastore_test(&logctx.log, &db, Uuid::new_v4()).await; let blueprint = BlueprintBuilder::build_empty_with_sleds( iter::once(SledUuid::new_v4()),