wiring up test database

Author: Jamesbarford

.github/workflows/ci.yml

@@ -33,7 +33,7 @@ jobs:
           key: linux
 
       - name: Run unit tests
-        run: cargo test --all
+        run: make test
 
       - name: Lint check
         run: cargo clippy --all -- -D warnings

Cargo.lock

@@ -0,0 +1,13 @@
+.PHONY: test start-postgres
+
+# Helper to spin up postgres
+start-postgres:
+	docker compose up -d pg_test
+
+# Helper to stop postgres & destroy the volume (ensures a clean rebuild on `start`)
+stop-postgres:
+	docker compose down -v
+
+# Run the tests with the database purring along in the background
+test: start-postgres
+	TEST_DB_URL="postgres://postgres:testpass@localhost/postgres" cargo test

Makefile

@@ -26,3 +26,4 @@ csv = "1"
 x509-cert = { version = "0.2.5", features = ["pem"] }
 
 intern = { path = "../intern" }
+uuid = "1.16.0"

database/Cargo.toml

@@ -14,6 +14,9 @@ pub mod metric;
 pub mod pool;
 pub mod selector;
 
+#[cfg(test)]
+mod tests;
+
 pub use pool::{Connection, Pool};
 
 intern!(pub struct Metric);

database/src/lib.rs

@@ -294,3 +294,92 @@ impl Pool {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use chrono::Utc;
+    use std::str::FromStr;
+
+    use super::*;
+    use crate::{tests::run_db_test, Commit, CommitType, Date};
+
+    /// Create a Commit
+    fn create_commit(commit_sha: &str, time: chrono::DateTime<Utc>, r#type: CommitType) -> Commit {
+        Commit {
+            sha: commit_sha.into(),
+            date: Date(time),
+            r#type,
+        }
+    }
+
+    #[tokio::test]
+    async fn pstat_returns_empty_vector_when_empty() {
+        run_db_test(|ctx| async {
+            // This is essentially testing the database testing framework is
+            // wired up correctly. Though makes sense that there should be
+            // an empty vector returned if there are no pstats.
+            let db = ctx.db_client();
+            let result = db.connection().await.get_pstats(&vec![], &vec![]).await;
+            let expected: Vec<Vec<Option<f64>>> = vec![];
+
+            assert_eq!(result, expected);
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn artifact_storage() {
+        run_db_test(|ctx| async {
+            let db = ctx.db_client();
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+
+            let artifact_one = ArtifactId::from(create_commit("abc", time, CommitType::Master));
+            let artifact_two = ArtifactId::Tag("nightly-2025-05-14".to_string());
+
+            let artifact_one_id_number = db.connection().await.artifact_id(&artifact_one).await;
+            let artifact_two_id_number = db.connection().await.artifact_id(&artifact_two).await;
+
+            // We cannot arbitrarily add random sizes to the artifact size
+            // table, as their is a constraint that the artifact must actually
+            // exist before attaching something to it.
+
+            // Artifact one inserts
+            db.connection()
+                .await
+                .record_artifact_size(artifact_one_id_number, "llvm.so", 32)
+                .await;
+            db.connection()
+                .await
+                .record_artifact_size(artifact_one_id_number, "llvm.a", 64)
+                .await;
+
+            // Artifact two inserts
+            db.connection()
+                .await
+                .record_artifact_size(artifact_two_id_number, "another-llvm.a", 128)
+                .await;
+
+            let result_one = db
+                .connection()
+                .await
+                .get_artifact_size(artifact_one_id_number)
+                .await;
+            let result_two = db
+                .connection()
+                .await
+                .get_artifact_size(artifact_two_id_number)
+                .await;
+
+            // artifact one
+            assert_eq!(Some(32 as u64), result_one.get("llvm.so").copied());
+            assert_eq!(Some(64 as u64), result_one.get("llvm.a").copied());
+            assert_eq!(None, result_one.get("another-llvm.a").copied());
+
+            // artifact two
+            assert_eq!(Some(128), result_two.get("another-llvm.a").copied());
+            Ok(ctx)
+        })
+        .await;
+    }
+}

database/src/pool.rs

@@ -25,7 +25,7 @@ impl Postgres {
 
 const CERT_URL: &str = "https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem";
 
-async fn make_client(db_url: &str) -> anyhow::Result<tokio_postgres::Client> {
+pub async fn make_client(db_url: &str) -> anyhow::Result<tokio_postgres::Client> {
     if db_url.contains("rds.amazonaws.com") {
         let mut builder = TlsConnector::builder();
         for cert in make_certificates().await {

database/src/pool/postgres.rs

@@ -0,0 +1,90 @@
+use std::future::Future;
+use tokio_postgres::config::Host;
+use tokio_postgres::Config;
+
+use crate::pool::postgres::make_client;
+use crate::Pool;
+
+/// Represents a connection to a Postgres database that can be
+/// used in integration tests to test logic that interacts with
+/// a database.
+pub(crate) struct TestContext {
+    db_name: String,
+    original_db_url: String,
+    // Pre-cached client to avoid creating unnecessary connections in tests
+    client: Pool,
+}
+
+impl TestContext {
+    async fn new(db_url: &str) -> Self {
+        let config: Config = db_url.parse().expect("Cannot parse connection string");
+
+        // Create a new database that will be used for this specific test
+        let client = make_client(&db_url)
+            .await
+            .expect("Cannot connect to database");
+        let db_name = format!("db{}", uuid::Uuid::new_v4().to_string().replace("-", ""));
+        client
+            .execute(&format!("CREATE DATABASE {db_name}"), &[])
+            .await
+            .expect("Cannot create database");
+        drop(client);
+
+        // We need to connect to the database against, because Postgres doesn't allow
+        // changing the active database mid-connection.
+        // There does not seem to be a way to turn the config back into a connection
+        // string, so construct it manually.
+        let test_db_url = format!(
+            "postgresql://{}:{}@{}:{}/{}",
+            config.get_user().unwrap(),
+            String::from_utf8(config.get_password().unwrap().to_vec()).unwrap(),
+            match &config.get_hosts()[0] {
+                Host::Tcp(host) => host,
+                Host::Unix(_) =>
+                    panic!("Unix sockets in Postgres connection string are not supported"),
+            },
+            &config.get_ports()[0],
+            db_name
+        );
+        let pool = Pool::open(test_db_url.as_str());
+
+        Self {
+            db_name,
+            original_db_url: db_url.to_string(),
+            client: pool,
+        }
+    }
+
+    pub(crate) fn db_client(&self) -> &Pool {
+        &self.client
+    }
+
+    async fn finish(self) {
+        // Cleanup the test database
+        // First, we need to stop using the database
+        drop(self.client);
+
+        // Then we need to connect to the default database and drop our test DB
+        let client = make_client(&self.original_db_url)
+            .await
+            .expect("Cannot connect to database");
+        client
+            .execute(&format!("DROP DATABASE {}", self.db_name), &[])
+            .await
+            .unwrap();
+    }
+}
+
+pub(crate) async fn run_db_test<F, Fut>(f: F)
+where
+    F: FnOnce(TestContext) -> Fut,
+    Fut: Future<Output = anyhow::Result<TestContext>>,
+{
+    if let Ok(db_url) = std::env::var("TEST_DB_URL") {
+        let ctx = TestContext::new(&db_url).await;
+        let ctx = f(ctx).await.expect("Test failed");
+        ctx.finish().await;
+    } else {
+        eprintln!("Skipping test because TEST_DB_URL was not passed");
+    }
+}

database/src/tests/mod.rs

@@ -0,0 +1,17 @@
+services:
+  pg_test:
+    image: postgres:16-alpine
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: testpass
+      POSTGRES_DB: postgres
+    ports:
+      - "5432:5432"                    # expose to host for tests
+    healthcheck:                       # wait until the DB is ready
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 2s
+      timeout: 2s
+      retries: 15
+    tmpfs:
+      - /var/lib/postgresql/data       # store data in RAM
+