Merge pull request #2119 from Jamesbarford/infra/add-database-testing-capability

Infra - wire up database for test

Author: Kobzol

.github/workflows/ci.yml

@@ -14,6 +14,17 @@ jobs:
   test_and_deploy:
     name: Test and deploy
     runs-on: ubuntu-22.04
+    env:
+      TEST_DB_URL: postgres://postgres:testpass@localhost:5432/postgres
+    services:
+      postgres:
+        image: postgres:16-alpine
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: testpass
+          POSTGRES_DB: postgres
+        ports:
+          - 5432:5432
     steps:
       - name: Checkout the source code
         uses: actions/checkout@v4

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

@@ -25,7 +25,9 @@ path = [
     "Cargo.lock",
     "Cargo.toml",
     "Dockerfile",
+    "docker-compose.yml",
     "LICENSE.md",
+    "Makefile",
     "README.md",
     "REUSE.toml",
     ".gitignore",

REUSE.toml

@@ -556,3 +556,22 @@ compilation of the final/leaf crate. Cargo only passes arguments after `--` to t
 therefore this does not affect the compilation of dependencies.
 2) Profiling/benchmarking - `cargo` is invoked with `--wrap-rustc-with <TOOL>`, which executes the
 specified profiling tool by `rustc-fake`.
+
+## How to test
+Run `make test`; in the root of the project there is a `Makefile` which
+presently exists to spin up/down a Postgres database, from a
+`docker-compose.yml`, then run `cargo test` with a `TEST_DB_URL` set. In
+concrete terms `make test` is a convenience for running;
+
+```bash
+docker compose up -d pg_test && \
+    TEST_DB_URL="postgres://postgres:testpass@localhost/postgres" cargo test
+```
+
+The above becomes cumbersome to type and easy to forget both how to setup the
+database and set the `TEST_DB_URL` environment variable.
+
+**Note: Windows**
+The tests for the database are disabled and will skip. This is due, at the time
+of writing (May 2025), to limitations with the GitHub Ci runner not supporting
+docker, hence unable to start the database for the tests.

collector/README.md

@@ -26,3 +26,6 @@ csv = "1"
 x509-cert = { version = "0.2.5", features = ["pem"] }
 
 intern = { path = "../intern" }
+
+[dev-dependencies]
+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,80 @@ 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 there is a constraint that the artifact must actually
+            // exist before attaching something to it.
+
+            let db = db.connection().await;
+
+            // Artifact one inserts
+            db.record_artifact_size(artifact_one_id_number, "llvm.so", 32)
+                .await;
+            db.record_artifact_size(artifact_one_id_number, "llvm.a", 64)
+                .await;
+
+            // Artifact two inserts
+            db.record_artifact_size(artifact_two_id_number, "another-llvm.a", 128)
+                .await;
+
+            let result_one = db.get_artifact_size(artifact_one_id_number).await;
+            let result_two = db.get_artifact_size(artifact_two_id_number).await;
+
+            // artifact one
+            assert_eq!(Some(32u64), result_one.get("llvm.so").copied());
+            assert_eq!(Some(64u64), 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,112 @@
+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);
+
+        let host = match &config
+            .get_hosts()
+            .first()
+            .expect("connection string must contain at least one host")
+        {
+            Host::Tcp(host) => host,
+
+            // This variant only exists on Unix targets, so the arm itself is
+            // cfg-gated to keep non-unix builds happy.
+            #[cfg(unix)]
+            Host::Unix(_) => panic!("Unix sockets in Postgres connection string are not supported"),
+
+            // On non-unix targets the enum has no other variants.
+            #[cfg(not(unix))]
+            _ => unreachable!("non-TCP hosts cannot appear on this platform"),
+        };
+
+        // 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(),
+            host,
+            &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 {
+        // The github CI does not yet support running containers on Windows,
+        // meaning that the test suite would fail.
+        if cfg!(unix) {
+            panic!("Aborting; `TEST_DB_URL` was not passed");
+        } else {
+            eprintln!(
+                "Skipping database test on platform {} `TEST_DB_URL` was not passed",
+                std::env::consts::OS
+            );
+        }
+    }
+}