Testing axum - Database Teardown
Rust
I love working with Rust.
Especially when writing backend systems, the guarantees that Rust provides are amazing. I can find errors at compile time that I would otherwise miss, I can refactor code I have not seen in a long time with confidence since the type system will catch my mistakes and point them out to me. I can be assured that I do not randomly break functionality on parts of the application because I misuse a function or input an invalid parameter. Things that would happen to me a lot in other languages without a strict type system.
Axum, Tower and SQLx
When writing a backend service, lately I have been using the Axum framework a lot. It is very easy to provide a production ready API with it and it integrates very well with Tower which has a lot of useful middlewares and utilities for building robust services. For database access, I have been using sqlx and I am quite amazed by it. It literally checks your SQL queries at compile time against your local database and raises an error if there is a mismatch. How great is that!
Missing the Ex
But you know, sometimes, as much as you love a language and a framework, you miss the comfort of things you were used to in other frameworks. Coming from Django and Pytest, I can say that the ecosystem there is very well developed and the comfort of developing and testing is very high. When running tests with pytest-django, automatically a database is created for the test run and teardown is also handled automatically, so you do not have to worry about it at all.
But what about the Rust Ecosystem? As a guide how to write production backend systems in Rust, I liked a lot the Zero to Production book. It extensively covers how to write a production ready backend system (but using the Actix Framework) and also covers testing. However, I found myself struggling with the database teardown part. Each time I run an integration test, I spawn the app in a separate thread and then run my test against it. In order to do that, I first must create and configure the database.
Usually my test would look something like this:
#[tokio::test]
async fn healthcheck() {
let app = spawn_app().await;
let response = app
.unauthenticated_client
.get(format!(
"{}{}",
app.address,
concrete_path(RouteResolver::Healthcheck.path())
))
.send()
.await
.expect("Failed to execute request");
assert_eq!(
response.status(),
200,
"/healthcheck must stay reachable"
);
app.teardown().await;
}
As you can see I already have a lot of helpers to slim this test down, but two are interesting.
- a
spawn_appfunction to start the application in a separate thread and return a handle to it. - a
teardownfunction to clean up the database and other resources after the test is done.
Writing a Teardown Function
Initially I was thinking about setting up one database for the lifecycle of the app and then just implement the teardown function via the Drop trait on the app handle, but I found that it was not working as expected. The Drop trait is called when the value goes out of scope (which is what we want), but there is a problem. You cannot await inside the Drop implementation since it is not an async function. This means that if you want to perform any asynchronous operations (like dropping a database), you cannot do it unless it runs synchronously. However, sqlx, the library which I dearly recommend when working with SQL databases in Rust, does not have a synchronous API.
So how could I automatically run the teardown code after the test is done? The solution that next came to my mind was to create a macro. While I try to use macros as little as possible, since they can make the code harder to read and reason about, sometimes they can be a powerful ally. So what are our requirements for this macro?
- Before the body of the test, it should run the
spawn_appand bind the result to a variable, let's sayapp. - After the body of the test, it should run the
teardownfunction on theappvariable. - It should be compatible with the
#[tokio::test]attribute, which we currently need to run our tests.
Why do we need this tokio macro in the first place? Because we need to asyncronously spawn the app and run asynchronous tests against it.
So here is the macro I came up with:
use proc_macro::TokenStream;
use quote::quote;
use syn::{ItemFn, Stmt, parse_macro_input, parse_quote};
/// Attribute macro for backend integration tests.
///
/// It turns a plain async test function into a Tokio test that:
/// 1) boots a fresh `TestApp` before the user test body runs, and
/// 2) always tears the app/database down afterwards.
#[proc_macro_attribute]
pub fn integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
// Parse the annotated item as a function so we can safely rewrite its AST.
let mut function = parse_macro_input!(item as ItemFn);
// This macro injects `.await` calls (`spawn_app` and `teardown`),
// therefore it only works on async functions. Thus, we check for asyncness early.
if function.sig.asyncness.is_none() {
return syn::Error::new_spanned(
function.sig.fn_token,
"`#[integration_test]` can only be used on async functions",
)
.to_compile_error()
.into();
}
// Normalize test attributes:
// - remove any existing `#[test]` / `#[tokio::test]`
// - add exactly one `#[tokio::test]`
// This avoids duplicate test attributes and keeps behavior predictable.
function.attrs.retain(|attr| {
let path = attr.path();
!(path.is_ident("test")
|| (path.segments.len() == 2
&& path.segments[0].ident == "tokio"
&& path.segments[1].ident == "test"))
});
function.attrs.insert(0, parse_quote!(#[tokio::test]));
// If the function body ends with a trailing expression (no semicolon),
// appending extra statements would make the generated code invalid.
// We turn that final expression into a terminated statement first.
if let Some(Stmt::Expr(_, semi)) = function.block.stmts.last_mut()
&& semi.is_none()
{
*semi = Some(Default::default());
}
// Inject test app lifecycle management around the original body.
// We use a crate-qualified path so the macro works from all test modules
// without requiring local `use spawn_app` imports.
function.block.stmts.insert(
0,
parse_quote! { let app = crate::helpers::application::spawn_app().await; },
);
function
.block
.stmts
// Teardown is appended last so test resources are released after assertions.
.push(parse_quote! { app.teardown().await; });
// Emit the rewritten function back to the compiler.
quote!(#function).into()
}
Remember, that there are different types of macros in Rust. For this type of macro (an attribute macro), we need to move it to a separate crate and import it in our test module.
When creating procedural macros, the definitions must reside in their own crate with a special crate type. This is for complex technical reasons that we hope to eliminate in the future. - The Rust Programming Language
This is a lot of code for such a small task, adding one line before and one line after the body of a test. So of course this is the moment where you expect me to say "but wait there is a much easier way". But unfortunately, I have yet to find it!
So here the test with the new macro:
#[integration_test]
async fn healthcheck() {
let response = app
.unauthenticated_client
.get(format!(
"{}{}",
app.address,
concrete_path(RouteResolver::Healthcheck.path())
))
.send()
.await
.expect("Failed to execute request");
assert_eq!(
response.status(),
200,
"/healthcheck must stay reachable"
);
}
Thanks for reading!