diff --git a/aggregator_core/src/datastore/test_util.rs b/aggregator_core/src/datastore/test_util.rs index 703989565..084058a73 100644 --- a/aggregator_core/src/datastore/test_util.rs +++ b/aggregator_core/src/datastore/test_util.rs @@ -27,7 +27,7 @@ use std::{ use testcontainers::{runners::AsyncRunner, ContainerRequest, ImageExt}; use tokio::{ io::{AsyncBufRead, AsyncBufReadExt}, - join, + join, spawn, sync::{ oneshot::{self, Sender}, Mutex, @@ -165,6 +165,7 @@ pub struct EphemeralDatastore { pool: Pool, datastore_key_bytes: Vec, migrator: Migrator, + db_name: String, } pub const TEST_DATASTORE_MAX_TRANSACTION_RETRIES: u64 = 1000; @@ -254,6 +255,21 @@ impl EphemeralDatastore { } } +impl Drop for EphemeralDatastore { + fn drop(&mut self) { + // Make a best-effort attempt to delete the database created for this datastore. This may + // fail if the container hosting the database server is deleted faster, which would make + // this moot. + let pool = self.pool.clone(); + let db_name = self.db_name.clone(); + spawn(async move { + if let Ok(conn) = pool.get().await { + let _ = conn.execute("DROP DATABASE $1", &[&db_name]).await; + } + }); + } +} + /// Builder to configure a new [`EphemeralDatastore`]. pub struct EphemeralDatastoreBuilder { schema_version: i64, @@ -383,6 +399,7 @@ impl EphemeralDatastoreBuilder { pool, datastore_key_bytes: generate_aead_key_bytes(), migrator, + db_name, } } }