move TestDatabase into database subcrate

Author: syphar

Cargo.lock

@@ -11,10 +11,19 @@ docs_rs_opentelemetry = { path = "../docs_rs_opentelemetry" }
 futures-util = { workspace = true }
 hex = "0.4.3"
 opentelemetry = { workspace = true }
+rand = { workspace = true, optional = true }
 sqlx = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 
+[dev-dependencies]
+rand = { workspace = true }
+docs_rs_opentelemetry = { path = "../docs_rs_opentelemetry", features = ["testing"] }
+test-case = { workspace = true }
+
 [features]
-testing = []
+testing = [
+    "dep:rand",
+    "docs_rs_opentelemetry/testing",
+]

crates/lib/docs_rs_database/Cargo.toml

@@ -3,6 +3,8 @@ mod errors;
 mod metrics;
 mod migrations;
 mod pool;
+#[cfg(any(test, feature = "testing"))]
+pub mod testing;
 
 pub use config::Config;
 pub use errors::PoolError;

crates/lib/docs_rs_database/src/lib.rs

@@ -30,7 +30,7 @@ impl Pool {
         Self::new_inner(config, DEFAULT_SCHEMA, otel_meter_provider).await
     }
 
-    #[cfg(feature = "testing")]
+    #[cfg(any(test, feature = "testing"))]
     pub async fn new_with_schema(
         config: &Config,
         schema: &str,

crates/lib/docs_rs_database/src/pool.rs

@@ -0,0 +1,3 @@
+mod test_env;
+
+pub use test_env::TestDatabase;

crates/lib/docs_rs_database/src/testing/mod.rs

@@ -0,0 +1,110 @@
+use crate::{AsyncPoolClient, Config, Pool, migrations};
+use anyhow::{Context as _, Result};
+use docs_rs_opentelemetry::AnyMeterProvider;
+use futures_util::TryStreamExt as _;
+use sqlx::Connection as _;
+use tokio::{runtime, task::block_in_place};
+use tracing::error;
+
+#[derive(Debug)]
+pub struct TestDatabase {
+    pool: Pool,
+    schema: String,
+    runtime: runtime::Handle,
+}
+
+impl TestDatabase {
+    pub async fn new(config: &Config, otel_meter_provider: &AnyMeterProvider) -> Result<Self> {
+        // A random schema name is generated and used for the current connection. This allows each
+        // test to create a fresh instance of the database to run within.
+        let schema = format!("docs_rs_test_schema_{}", rand::random::<u64>());
+
+        let pool = Pool::new_with_schema(config, &schema, otel_meter_provider).await?;
+
+        let mut conn = sqlx::PgConnection::connect(&config.database_url).await?;
+        sqlx::query(&format!("CREATE SCHEMA {schema}"))
+            .execute(&mut conn)
+            .await
+            .context("error creating schema")?;
+        sqlx::query(&format!("SET search_path TO {schema}, public"))
+            .execute(&mut conn)
+            .await
+            .context("error setting search path")?;
+        migrations::migrate(&mut conn, None)
+            .await
+            .context("error running migrations")?;
+
+        // Move all sequence start positions 10000 apart to avoid overlapping primary keys
+        let sequence_names: Vec<_> = sqlx::query!(
+            "SELECT relname
+             FROM pg_class
+             INNER JOIN pg_namespace ON
+                 pg_class.relnamespace = pg_namespace.oid
+             WHERE pg_class.relkind = 'S'
+                 AND pg_namespace.nspname = $1
+            ",
+            schema,
+        )
+        .fetch(&mut conn)
+        .map_ok(|row| row.relname)
+        .try_collect()
+        .await?;
+
+        for (i, sequence) in sequence_names.into_iter().enumerate() {
+            let offset = (i + 1) * 10000;
+            sqlx::query(&format!(
+                r#"ALTER SEQUENCE "{sequence}" RESTART WITH {offset};"#
+            ))
+            .execute(&mut conn)
+            .await?;
+        }
+
+        Ok(TestDatabase {
+            pool,
+            schema,
+            runtime: runtime::Handle::current(),
+        })
+    }
+
+    pub fn pool(&self) -> &Pool {
+        &self.pool
+    }
+
+    pub async fn async_conn(&self) -> AsyncPoolClient {
+        self.pool
+            .get_async()
+            .await
+            .expect("failed to get a connection out of the pool")
+    }
+}
+
+impl Drop for TestDatabase {
+    fn drop(&mut self) {
+        let pool = self.pool.clone();
+        let schema = self.schema.clone();
+        let runtime = self.runtime.clone();
+
+        block_in_place(move || {
+            runtime.block_on(async move {
+                let Ok(mut conn) = pool.get_async().await else {
+                    error!("error in drop impl");
+                    return;
+                };
+
+                let migration_result = migrations::migrate(&mut conn, Some(0)).await;
+
+                if let Err(e) = sqlx::query(format!("DROP SCHEMA {} CASCADE;", schema).as_str())
+                    .execute(&mut *conn)
+                    .await
+                {
+                    error!("failed to drop test schema {}: {}", schema, e);
+                    return;
+                }
+
+                if let Err(err) = migration_result {
+                    error!(?err, "error reverting migrations");
+                }
+            })
+        });
+    }
+}

crates/lib/docs_rs_database/src/testing/test_env.rs

@@ -1,4 +1,3 @@
-use super::TestDatabase;
 use crate::{
     db::{initialize_build, initialize_crate, initialize_release, update_build_status},
     docbuilder::{DocCoverage, RUSTDOC_JSON_COMPRESSION_ALGORITHMS},
@@ -8,6 +7,7 @@ use anyhow::{Context, bail};
 use base64::{Engine, engine::general_purpose::STANDARD as b64};
 use chrono::{DateTime, Utc};
 use docs_rs_cargo_metadata::{Dependency, MetadataPackage, Target};
+use docs_rs_database::testing::TestDatabase;
 use docs_rs_registry_api::{CrateData, CrateOwner, ReleaseData};
 use docs_rs_storage::{
     AsyncStorage, CompressionAlgorithm, FileEntry, RustdocJsonFormatVersion,

src/test/fakes.rs

@@ -12,28 +12,22 @@ use anyhow::{Context as _, anyhow};
 use axum::body::Bytes;
 use axum::{Router, body::Body, http::Request, response::Response as AxumResponse};
 use axum_extra::headers::{ETag, HeaderMapExt as _};
-use docs_rs_database::{AsyncPoolClient, Pool};
+use docs_rs_database::testing::TestDatabase;
 use docs_rs_fastly::Cdn;
 use docs_rs_headers::{IfNoneMatch, SURROGATE_CONTROL, SurrogateKeys};
-use docs_rs_opentelemetry::{
-    AnyMeterProvider,
-    testing::{CollectedMetrics, TestMetrics},
-};
+use docs_rs_opentelemetry::testing::{CollectedMetrics, TestMetrics};
 use docs_rs_storage::{AsyncStorage, Storage, StorageKind, testing::TestStorage};
 use docs_rs_types::Version;
 use fn_error_context::context;
-use futures_util::stream::TryStreamExt;
 use http::{
     HeaderMap, HeaderName, HeaderValue, StatusCode,
     header::{CACHE_CONTROL, CONTENT_TYPE},
 };
 use http_body_util::BodyExt;
 use serde::de::DeserializeOwned;
-use sqlx::Connection as _;
 use std::{collections::HashMap, fs, future::Future, panic, rc::Rc, sync::Arc};
-use tokio::{runtime, task::block_in_place};
+use tokio::runtime;
 use tower::ServiceExt;
-use tracing::error;
 
 // testing krate name constants
 pub(crate) const KRATE: &str = "krate";
@@ -452,8 +446,7 @@ impl TestEnvironment {
         fs::create_dir_all(config.registry_index_path.clone())?;
 
         let test_metrics = TestMetrics::new();
-
-        let test_db = TestDatabase::new(&config, test_metrics.provider())
+        let test_db = TestDatabase::new(&config.database, test_metrics.provider())
             .await
             .context("can't initialize test database")?;
 
@@ -542,108 +535,3 @@ impl TestEnvironment {
         fakes::FakeRelease::new(self.async_db(), self.context.async_storage.clone())
     }
 }
-
-#[derive(Debug)]
-pub(crate) struct TestDatabase {
-    pool: Pool,
-    schema: String,
-    runtime: runtime::Handle,
-}
-
-impl TestDatabase {
-    async fn new(config: &Config, otel_meter_provider: &AnyMeterProvider) -> Result<Self> {
-        // A random schema name is generated and used for the current connection. This allows each
-        // test to create a fresh instance of the database to run within.
-        let schema = format!("docs_rs_test_schema_{}", rand::random::<u64>());
-
-        let config = &config.database;
-
-        let pool = Pool::new_with_schema(config, &schema, otel_meter_provider).await?;
-
-        let mut conn = sqlx::PgConnection::connect(&config.database_url).await?;
-        sqlx::query(&format!("CREATE SCHEMA {schema}"))
-            .execute(&mut conn)
-            .await
-            .context("error creating schema")?;
-        sqlx::query(&format!("SET search_path TO {schema}, public"))
-            .execute(&mut conn)
-            .await
-            .context("error setting search path")?;
-        docs_rs_database::migrate(&mut conn, None)
-            .await
-            .context("error running migrations")?;
-
-        // Move all sequence start positions 10000 apart to avoid overlapping primary keys
-        let sequence_names: Vec<_> = sqlx::query!(
-            "SELECT relname
-             FROM pg_class
-             INNER JOIN pg_namespace ON
-                 pg_class.relnamespace = pg_namespace.oid
-             WHERE pg_class.relkind = 'S'
-                 AND pg_namespace.nspname = $1
-            ",
-            schema,
-        )
-        .fetch(&mut conn)
-        .map_ok(|row| row.relname)
-        .try_collect()
-        .await?;
-
-        for (i, sequence) in sequence_names.into_iter().enumerate() {
-            let offset = (i + 1) * 10000;
-            sqlx::query(&format!(
-                r#"ALTER SEQUENCE "{sequence}" RESTART WITH {offset};"#
-            ))
-            .execute(&mut conn)
-            .await?;
-        }
-
-        Ok(TestDatabase {
-            pool,
-            schema,
-            runtime: runtime::Handle::current(),
-        })
-    }
-
-    pub(crate) fn pool(&self) -> &Pool {
-        &self.pool
-    }
-
-    pub(crate) async fn async_conn(&self) -> AsyncPoolClient {
-        self.pool
-            .get_async()
-            .await
-            .expect("failed to get a connection out of the pool")
-    }
-}
-
-impl Drop for TestDatabase {
-    fn drop(&mut self) {
-        let pool = self.pool.clone();
-        let schema = self.schema.clone();
-        let runtime = self.runtime.clone();
-
-        block_in_place(move || {
-            runtime.block_on(async move {
-                let Ok(mut conn) = pool.get_async().await else {
-                    error!("error in drop impl");
-                    return;
-                };
-
-                let migration_result = docs_rs_database::migrate(&mut conn, Some(0)).await;
-
-                if let Err(e) = sqlx::query(format!("DROP SCHEMA {} CASCADE;", schema).as_str())
-                    .execute(&mut *conn)
-                    .await
-                {
-                    error!("failed to drop test schema {}: {}", schema, e);
-                    return;
-                }
-
-                if let Err(err) = migration_result {
-                    error!(?err, "error reverting migrations");
-                }
-            })
-        });
-    }
-}

src/test/mod.rs

@@ -723,10 +723,11 @@ mod tests {
     use super::*;
     use crate::db::update_build_status;
     use crate::test::{
-        AxumResponseTestExt, AxumRouterTestExt, FakeBuild, TestDatabase, TestEnvironment,
-        async_wrapper, fake_release_that_failed_before_build,
+        AxumResponseTestExt, AxumRouterTestExt, FakeBuild, TestEnvironment, async_wrapper,
+        fake_release_that_failed_before_build,
     };
     use anyhow::Error;
+    use docs_rs_database::testing::TestDatabase;
     use docs_rs_registry_api::CrateOwner;
     use docs_rs_types::KrateName;
     use http::StatusCode;

src/web/crate_details.rs

@@ -628,9 +628,9 @@ mod test {
     use super::*;
     use crate::docbuilder::DocCoverage;
     use crate::test::{
-        AxumResponseTestExt, AxumRouterTestExt, FakeBuild, TestDatabase, TestEnvironment,
-        async_wrapper,
+        AxumResponseTestExt, AxumRouterTestExt, FakeBuild, TestEnvironment, async_wrapper,
     };
+    use docs_rs_database::testing::TestDatabase;
     use docs_rs_types::ReleaseId;
     use kuchikiki::traits::TendrilSink;
     use pretty_assertions::assert_eq;