Merge pull request #1892 from Kobzol/db-tests

Add simple test infrastructure for testing DB queries

Author: jackh726

.github/workflows/ci.yml

@@ -18,6 +18,27 @@ jobs:
       run: rustup update stable && rustup default stable && rustup component add rustfmt
     - run: cargo fmt --all --check
 
+  test:
+    name: Test
+    runs-on: ubuntu-latest
+    env:
+      TEST_DB_URL: postgres://postgres:postgres@localhost:5432/postgres
+    services:
+      postgres:
+        image: postgres:14
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: postgres
+        ports:
+          - 5432:5432
+    steps:
+      - uses: actions/checkout@v4
+      - run: rustup toolchain install stable --profile minimal
+      - uses: Swatinem/rust-cache@v2
+      - name: Run tests
+        run: cargo test --workspace --all-targets
+
   ci:
     name: CI
     runs-on: ubuntu-latest

README.md

@@ -39,12 +39,12 @@ The general overview of what you will need to do:
 3. [Configure webhook forwarding](#configure-webhook-forwarding)
 4. Configure the `.env` file:
 
-   1. Copy `.env.sample` to `.env`
-   2. `GITHUB_TOKEN`: This is a token needed for Triagebot to send requests to GitHub. Go to GitHub Settings > Developer Settings > Personal Access Token, and create a new token. The `repo` permission should be sufficient.
-      If this is not set, Triagebot will also look in `~/.gitconfig` in the `github.oauth-token` setting.
-   3. `DATABASE_URL`: This is the URL to the database. See [Configuring a database](#configuring-a-database).
-   4. `GITHUB_WEBHOOK_SECRET`: Enter the secret you entered in the webhook above.
-   5. `RUST_LOG`: Set this to `debug`.
+    1. Copy `.env.sample` to `.env`
+    2. `GITHUB_TOKEN`: This is a token needed for Triagebot to send requests to GitHub. Go to GitHub Settings > Developer Settings > Personal Access Token, and create a new token. The `repo` permission should be sufficient.
+       If this is not set, Triagebot will also look in `~/.gitconfig` in the `github.oauth-token` setting.
+    3. `DATABASE_URL`: This is the URL to the database. See [Configuring a database](#configuring-a-database).
+    4. `GITHUB_WEBHOOK_SECRET`: Enter the secret you entered in the webhook above.
+    5. `RUST_LOG`: Set this to `debug`.
 
 5. Run `cargo run --bin triagebot`. This starts the http server listening for webhooks on port 8000.
 6. Add a `triagebot.toml` file to the main branch of your GitHub repo with whichever services you want to try out.
@@ -109,15 +109,28 @@ You need to sign up for a free account, and also deal with configuring the GitHu
 3. Configure GitHub webhooks in the test repo you created.
    In short:
 
-   1. Go to the settings page for your GitHub repo.
-   2. Go to the webhook section.
-   3. Click "Add webhook"
-   4. Include the settings:
+    1. Go to the settings page for your GitHub repo.
+    2. Go to the webhook section.
+    3. Click "Add webhook"
+    4. Include the settings:
 
-      * Payload URL: This is the URL to your Triagebot server, for example http://7e9ea9dc.ngrok.io/github-hook. This URL is displayed when you ran the `ngrok` command above.
-      * Content type: application/json
-      * Secret: Enter a shared secret (some longish random text)
-      * Events: "Send me everything"
+        * Payload URL: This is the URL to your Triagebot server, for example http://7e9ea9dc.ngrok.io/github-hook. This URL is displayed when you ran the `ngrok` command above.
+        * Content type: application/json
+        * Secret: Enter a shared secret (some longish random text)
+        * Events: "Send me everything"
+
+### Cargo tests
+
+You can run Cargo tests using `cargo test`. If you also want to run tests that access a Postgres database, you can specify an environment variables `TEST_DB_URL`, which should contain a connection string pointing to a running Postgres database instance:
+
+```bash
+$ docker run --rm -it -p5432:5432 \
+  -e POSTGRES_USER=triagebot \
+  -e POSTGRES_PASSWORD=triagebot \
+  -e POSTGRES_DB=triagebot \
+  postgres:14
+$ TEST_DB_URL=postgres://triagebot:triagebot@localhost:5432/triagebot cargo test
+```
 
 ## License
 

src/db.rs

@@ -23,6 +23,7 @@ static CERTIFICATE_PEMS: LazyLock<Vec<u8>> = LazyLock::new(|| {
 pub struct ClientPool {
     connections: Arc<Mutex<Vec<tokio_postgres::Client>>>,
     permits: Arc<Semaphore>,
+    db_url: String,
 }
 
 pub struct PooledClient {
@@ -54,10 +55,11 @@ impl std::ops::DerefMut for PooledClient {
 }
 
 impl ClientPool {
-    pub fn new() -> ClientPool {
+    pub fn new(db_url: String) -> ClientPool {
         ClientPool {
             connections: Arc::new(Mutex::new(Vec::with_capacity(16))),
             permits: Arc::new(Semaphore::new(16)),
+            db_url,
         }
     }
 
@@ -79,15 +81,14 @@ impl ClientPool {
         }
 
         PooledClient {
-            client: Some(make_client().await.unwrap()),
+            client: Some(make_client(&self.db_url).await.unwrap()),
             permit,
             pool: self.connections.clone(),
         }
     }
 }
 
-async fn make_client() -> anyhow::Result<tokio_postgres::Client> {
-    let db_url = std::env::var("DATABASE_URL").expect("needs DATABASE_URL");
+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() {
@@ -230,8 +231,9 @@ pub async fn schedule_job(
     Ok(())
 }
 
-pub async fn run_scheduled_jobs(ctx: &Context, db: &DbClient) -> anyhow::Result<()> {
-    let jobs = get_jobs_to_execute(&db).await.unwrap();
+pub async fn run_scheduled_jobs(ctx: &Context) -> anyhow::Result<()> {
+    let db = &ctx.db.get().await;
+    let jobs = get_jobs_to_execute(&db).await?;
     tracing::trace!("jobs to execute: {:#?}", jobs);
 
     for job in jobs.iter() {

src/handlers/assign.rs

@@ -39,6 +39,7 @@ use tracing as log;
 #[cfg(test)]
 mod tests {
     mod tests_candidates;
+    mod tests_db;
     mod tests_from_diff;
 }
 

src/handlers/assign/tests/tests_db.rs

@@ -0,0 +1,30 @@
+#[cfg(test)]
+mod tests {
+    use crate::handlers::assign::filter_by_capacity;
+    use crate::tests::run_test;
+    use std::collections::HashSet;
+
+    #[tokio::test]
+    async fn find_reviewers_no_review_prefs() {
+        run_test(|ctx| async move {
+            ctx.add_user("usr1", 1).await;
+            ctx.add_user("usr2", 1).await;
+            let _users =
+                filter_by_capacity(ctx.db_client(), &candidates(&["usr1", "usr2"])).await?;
+            // FIXME: this test fails, because the query is wrong
+            // check_users(users, &["usr1", "usr2"]);
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    fn candidates(users: &[&'static str]) -> HashSet<&'static str> {
+        users.into_iter().copied().collect()
+    }
+
+    fn check_users(users: HashSet<String>, expected: &[&'static str]) {
+        let mut users: Vec<String> = users.into_iter().collect();
+        users.sort();
+        assert_eq!(users, expected);
+    }
+}

src/lib.rs

@@ -26,6 +26,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/main.rs

@@ -242,7 +242,8 @@ async fn serve_req(
 }
 
 async fn run_server(addr: SocketAddr) -> anyhow::Result<()> {
-    let pool = db::ClientPool::new();
+    let db_url = std::env::var("DATABASE_URL").expect("needs DATABASE_URL");
+    let pool = db::ClientPool::new(db_url.clone());
     db::run_migrations(&mut *pool.get().await)
         .await
         .context("database migrations")?;
@@ -271,7 +272,7 @@ async fn run_server(addr: SocketAddr) -> anyhow::Result<()> {
 
     // Run all jobs that have a schedule (recurring jobs)
     if !is_scheduled_jobs_disabled() {
-        spawn_job_scheduler();
+        spawn_job_scheduler(db_url);
         spawn_job_runner(ctx.clone());
     }
 
@@ -361,11 +362,12 @@ async fn spawn_job_oneoffs(ctx: Arc<Context>) {
 /// The scheduler wakes up every `JOB_SCHEDULING_CADENCE_IN_SECS` seconds to
 /// check if there are any jobs ready to run. Jobs get inserted into the the
 /// database which acts as a queue.
-fn spawn_job_scheduler() {
+fn spawn_job_scheduler(db_url: String) {
     task::spawn(async move {
         loop {
+            let db_url = db_url.clone();
             let res = task::spawn(async move {
-                let pool = db::ClientPool::new();
+                let pool = db::ClientPool::new(db_url);
                 let mut interval =
                     time::interval(time::Duration::from_secs(JOB_SCHEDULING_CADENCE_IN_SECS));
 
@@ -401,13 +403,12 @@ fn spawn_job_runner(ctx: Arc<Context>) {
         loop {
             let ctx = ctx.clone();
             let res = task::spawn(async move {
-                let pool = db::ClientPool::new();
                 let mut interval =
                     time::interval(time::Duration::from_secs(JOB_PROCESSING_CADENCE_IN_SECS));
 
                 loop {
                     interval.tick().await;
-                    db::run_scheduled_jobs(&ctx, &*pool.get().await)
+                    db::run_scheduled_jobs(&ctx)
                         .await
                         .context("run database scheduled jobs")
                         .unwrap();

src/tests.rs

@@ -0,0 +1,93 @@
+use crate::db;
+use crate::db::make_client;
+use crate::db::notifications::record_username;
+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,
+        }
+    }
+
+    pub fn db_client(&self) -> &tokio_postgres::Client {
+        &self.client
+    }
+
+    pub async fn add_user(&self, name: &str, id: u64) {
+        record_username(&self.client, id, name)
+            .await
+            .expect("Cannot create user");
+    }
+
+    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");
+    }
+}