Add helper code for running opt-in tests with a DB

Author: Kobzol

src/db.rs

@@ -93,7 +93,7 @@ impl ClientPool {
     }
 }
 
-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() {

src/lib.rs

@@ -29,6 +29,9 @@ mod team_data;
 pub mod triage;
 pub mod zulip;
 
+#[cfg(test)]
+mod tests;
+
 /// The name of a webhook event.
 #[derive(Debug)]
 pub enum EventName {

src/tests.rs

@@ -0,0 +1,82 @@
+use crate::db;
+use crate::db::make_client;
+use std::future::Future;
+use tokio_postgres::Config;
+
+/// Represents a connection to a Postgres database that can be
+/// used in integration tests to test logic that interacts with
+/// a database.
+pub struct TestContext {
+    client: tokio_postgres::Client,
+    db_name: String,
+    original_db_url: String,
+    conn_handle: tokio::task::JoinHandle<()>,
+}
+
+impl TestContext {
+    async fn new(db_url: &str) -> Self {
+        let mut 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.
+        config.dbname(&db_name);
+        let (mut client, connection) = config
+            .connect(tokio_postgres::NoTls)
+            .await
+            .expect("Cannot connect to the newly created database");
+        let conn_handle = tokio::spawn(async move {
+            connection.await.unwrap();
+        });
+
+        db::run_migrations(&mut client)
+            .await
+            .expect("Cannot run database migrations");
+        Self {
+            client,
+            db_name,
+            original_db_url: db_url.to_string(),
+            conn_handle,
+        }
+    }
+
+    async fn finish(self) {
+        // Cleanup the test database
+        // First, we need to stop using the database
+        drop(self.client);
+        self.conn_handle.await.unwrap();
+
+        // 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 async fn run_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");
+    }
+}