Concurrent Unit Testing in Rocket with a Database

I’ve been writing a webapp using Rocket lately, and as part of that, I wanted to write unit tests for my routes. This requires testing against a database. I came up with a solution to enable concurrent unit testing when testing a Rocket application, which vastly improves the runtime of cargo test.

Rocket testing

Rocket provides unit testing utilities that allow you to launch a virtual “server” of your application that isn’t actually exposed to the network. This means that your unit tests don’t expose any network resources and aren’t competing for a port assignment. They otherwise behave as if they were an actual running instance of your application, making them perfect for testing routes.

The problem

There is a problem with testing routes: if the routes you’re testing interact with a database at all, you need to somehow isolate each test’s database from each other. Otherwise, your tests will see the state of other tests, which will cause bizarre failures. Rocket’s TODO example gets around this issue by wrapping the tests in a macro that clears out the database before each test. This has some problems, however.

Firstly, since all the tests use the same database, you have to force each test to run in serial via a global mutex; otherwise, they’ll observe each other’s state when the tests interleave database operations. For an application with a lot of tests, this slows down testing a lot. The only thing requiring these tests to run in sequence is their common database.

The second problem is a code maintainability problem. The TODO example relies on someone writing code to manually clear out the entire database before the test runs. If you add another table to the database, you have to write code to clear that table out before each test as well. If you forget, state leaks across tests. This is a case where the default behavior is to cause hard-to-debug test failures, and it’s not good.

My situation

In my case, my application is limited to one Postgres database. If I could instruct the test to use a different database, then I could isolate each test from each other. I don’t have to write code to clear out the database each time, since it’s unique to that test. Postgres is totally capable of handling multiple tests interacting with different databases.

If your application uses SQLite instead, it’s even easier - you don’t have a database server to interact with; you can just open an in-memory SQLite database for each test. If you use a different database system, it’s still easy - you create a new temporary database for each unit test invocation.

Assigning each test a database

I want to write a function that creates a Client for each test. This function should give me a Client with a unique database.

use chrono::Utc;
use diesel::prelude::*;
use figment::{providers::Serialized, Figment};
use rocket_sync_db_pools::Config;
use rocket::local::asynchronous::Client;

pub(crate) async fn test_client(test_config: crate::config::Config) -> Client {
}

Foreword: Test environment

Before we get into the implementation of test_client, I want to make some information about the environment in which these tests are running apparent. Firstly, they are running in a Docker container. Secondly, there is a Postgres server running in another container that is accessible via the database address on port 5432. When the tests finish running, the database container is shut down and all data is deleted.

Picking a database name

The database name needs to be unique for all tests, but it doesn’t need to be globally unique, because the database server gets torn down. However, since the application already uses the uuid crate, I opted to use a combination of UNIX timestamp + version 4 UUID to determine the name of each test’s database.

let timestamp = chrono::Utc::now().timestamp();
let uuid = uuid::Uuid::new_v4();
let db_name = format!(
    "recipemgr_{}_{}",
    timestamp,
    uuid.to_simple().to_string());

Creating the database

My application uses Diesel as its ORM and database driver, so this code is Diesel-specific, but it’s fairly straightforward to translate to another database driver, since it’s all plain SQL. We need to create the database we’re telling the application to use:

let conn = PgConnection::establish("postgres://username:password@database:5432")
    .expect("Unable to connect to database server");

let query = diesel::sql_query(format!("CREATE DATABASE {};", db_name));

query.execute(&conn).expect("Unable to create database");

let query = diesel::sql_query(format!(
    "GRANT ALL PRIVILEGES ON DATABASE {} TO username;",
    db_name
));

query
    .execute(&conn)
    .expect("Unable to grant privileges to test user");

Note: No awaits are necessary because Diesel is a synchronous API.

Once this code is done, we’ve created a new database with the given database name, empty for our tests to use. We’ll also need to run the database migrations to set up the database’s schema. These are helpfully embedded in the application via the diesel_migrations crate. The one snag I ran into is that we need to re-establish our PgConnection to scope it to the database we just created, but this is very simple to do.

// This will drop the old connection.
let conn = PgConnection::establish(&format!(
        "postgres://username:password@database:5432/{}",
        db_name
    ))
    .expect("Unable to connect to database server");

crate::embedded_migrations::run(&conn).expect("Couldn't run migrations");

Overriding configuration

My application uses Rocket’s rocket_sync_db_pools crate to connect to its Postgres database. Which database the application connects to is configurable, and this configuration can be modified by arbitrary code before the Rocket instance is started. In order to assign each test a unique database, I need to create the Rocket instance with a custom configuration provider that overrides part of the configuration defaults. This was also a good opportunity to make configuration settings testable.

In order to override the database configuration, I need to create a custom Figment instance, which is the library that Rocket uses for handling configuration, and override a specific value in it. I want the merge method of Figment, which will replace values in the configuration, rather than join, which does nothing if the value is already present.

The database name is a combination of the current timestamp and a UUID. This is sufficient information that two tests getting assigned the same database name should never happen, or happen so rarely that it isn’t worth guarding against.

let db_config = rocket_sync_db_pools::Config {
    // `database` is the name of the Docker container that the Postgres
    // server is running in.
    url: format!("postgres://username:password@database:5432/{}", db_name),
    // Using a smaller pool size prevents test failures when many tests
    // are invoked concurrently. The default pool size of n_workers * 4 will
    // quickly consume all available connections on the Postgres server.
    pool_size: 5,
    // A lower timeout will improve test performance if a connection to the
    // database can't be made.
    timeout: 5,
};

let figment = Figment::from(rocket::Config::default())
    // Merge in the application's default configuration.
    .merge(Serialized::defaults(crate::config::Config::default()))
    // Override the application's default configuration with whatever
    // configuration the test requires.
    .merge(Serialized::defaults(test_config))
    // Override the database configuration for the application's database
    // with our custom database configuration.
    .merge(Serialized::global("databases.recipemgr_db", db_config));

// Create a Rocket instance using our custom configuration provider.
let rocket = rocket::custom(figment);

// ...more happens to configure the Rocket...

Client::tracked(rocket)
    .await
    .expect("Couldn't create client")

Cleaning up

It is important to note that if you use this technique, you should be sure to delete the test databases after the tests are finished, lest you accumulate databases endlessly. I don’t do this here, because the entire database server will be cleaned up at the end of the test invocation, but if your database server persists beyond the lifetime of your test executable, you’ll want to manually execute DROP DATABASE db_name SQL commands at the end of each test. A good way to automate this would be to return a struct from test_client that has a Drop implementation.

Final code

use chrono::Utc;
use diesel::prelude::*;
use figment::{providers::Serialized, Figment};
use rocket_sync_db_pools::Config;
use rocket::local::asynchronous::Client;

pub(crate) async fn test_client(test_config: crate::config::Config) -> Client {
    let timestamp = Utc::now().timestamp();
    let uuid = uuid::Uuid::new_v4();
    let db_name = format!("recipemgr_{}_{}", timestamp, uuid.to_simple().to_string());

    {
        let conn = PgConnection::establish("postgres://username:password@database:5432")
            .expect("Unable to connect to database server");

        let query = diesel::sql_query(format!("CREATE DATABASE {};", db_name));

        query.execute(&conn).expect("Unable to create database");

        let query = diesel::sql_query(format!(
            "GRANT ALL PRIVILEGES ON DATABASE {} TO username;",
            db_name
        ));

        query
            .execute(&conn)
            .expect("Unable to grant privileges to test user");

        let conn = PgConnection::establish(&format!(
            "postgres://username:password@database:5432/{}",
            db_name
        ))
        .expect("Unable to connect to database server");
        crate::embedded_migrations::run(&conn).expect("Couldn't run migrations");
    }

    let db_config = rocket_sync_db_pools::Config {
        url: format!("postgres://username:password@database:5432/{}", db_name),
        pool_size: 5,
        timeout: 5,
    };

    let figment = Figment::from(rocket::Config::default())
        .merge(Serialized::defaults(crate::config::Config::default()))
        .merge(Serialized::defaults(test_config))
        .merge(Serialized::global("databases.recipemgr_db", db_config));
    
    let rocket = rocket::custom(figment);

    // ...more happens to configure the Rocket...

    Client::tracked(rocket)
        .await
        .expect("Couldn't create client")
}