How to design an Axum server with test-friendliness in mind?

Question

When I tried to build an application with axum, I failed to separate the framework from my handler. With Go, the classic way is define an Interface, implement it and register the handler to framework. In this way, it's easy to provide a mock handler to test with. However, I couldn't make it work with Axum. I defined a trait just like above, but it wouldn't compile:

use std::net::ToSocketAddrs;
use std::sync::{Arc, Mutex};
use serde_derive::{Serialize, Deserialize};
use serde_json::json;
use axum::{Server, Router, Json};

use axum::extract::Extension;
use axum::routing::BoxRoute;
use axum::handler::get;

#[tokio::main]
async fn main() {
    let app = new_router(
        Foo{}
    );
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

trait Handler {
    fn get(&self, get: GetRequest) -> Result<GetResponse, String>;
}

struct Foo {}

impl Handler for Foo {
    fn get(&self, req: GetRequest) -> Result<GetResponse, String> {
        Ok(GetResponse{ message: "It works.".to_owned()})
    }
}

fn new_router<T:Handler>(handler: T) -> Router<BoxRoute> {
    Router::new()
        .route("/", get(helper))
        .boxed()
}

fn helper<T:Handler>(
    Extension(mut handler): Extension<T>,
    Json(req): Json<GetRequest>
) -> Json<GetResponse> {
    Json(handler.get(req).unwrap())
}

#[derive(Debug, Serialize, Deserialize)]
struct GetRequest {
    // omited
}

#[derive(Debug, Serialize, Deserialize)]
struct GetResponse {
    message: String
    // omited
}

error[E0599]: the method `boxed` exists for struct `Router<axum::routing::Layered<Trace<axum::routing::Layered<AddExtension<Nested<Router<BoxRoute>, Route<axum::handler::OnMethod<fn() -> impl Future {direct}, _, (), EmptyRouter>, EmptyRouter<_>>>, T>>, SharedClassifier<ServerErrorsAsFailures>>>>`, but its trait bounds were not satisfied
   --> src/router.rs:25:10
    |
25  |         .boxed()
    |          ^^^^^ method cannot be called on `Router<axum::routing::Layered<Trace<axum::routing::Layered<AddExtension<Nested<Router<BoxRoute>, Route<axum::handler::OnMethod<fn() -> impl Future {direct}, _, (), EmptyRouter>, EmptyRouter<_>>>, T>>, SharedClassifier<ServerErrorsAsFailures>>>>` due to unsatisfied trait bounds
    | 
   ::: /Users/lebrancebw/.cargo/registry/src/github.com-1ecc6299db9ec823/axum-0.2.5/src/routing/mod.rs:876:1
    |
876 | pub struct Layered<S> {
    | --------------------- doesn't satisfy `<_ as tower_service::Service<Request<_>>>::Error = _`
    |
    = note: the following trait bounds were not satisfied:
            `<axum::routing::Layered<Trace<axum::routing::Layered<AddExtension<Nested<Router<BoxRoute>, Route<axum::handler::OnMethod<fn() -> impl Future {direct}, _, (), EmptyRouter>, EmptyRouter<_>>>, T>>, SharedClassifier<ServerErrorsAsFailures>>> as tower_service::Service<Request<_>>>::Error = _`

I guess the key point is my design is apparently not "rustic". Is there a way to structure an Axum project that lends itself to testing easily?

iTollu · Accepted Answer · 2022-01-09T15:27:50.103

The question is what you want to test. I will assume that you have some core logic and an HTTP layer. And you want to make sure that:

requests are routed correctly;
requests are parsed correctly;
the core logic is called with the expected parameters;
and return values from the core are correctly formatted into HTTP responses.

To test it you want to spawn an instance of the server with the core logic mocked out. @lukemathwalker in his blog and book "Zero To Production In Rust" has a very nice description of how to spawn an app for testing through the actual TCP port. It is written for Actix-Web, but the idea applies to Axum as well.

You should not use axum::Server::bind, but rather use axum::Server::from_tcp to pass it a std::net::TcpListner which allows you to spawn a test server on any available port using `TcpListener::bind("127.0.0.1:0").

To make the core logic injectable (and mockable) I declare it as a struct and implement all the business methods on it. Like this:

pub struct Core {
    public_url: Url,
    secret_key: String,
    storage: Storage,
    client: SomeThirdPartyClient,
}

impl Core {
    pub async fn create_something(
        &self,
        new_something: NewSomething,
    ) -> Result<Something, BusinessErr> {
    ...
}

With all these pieces you can write a function to start the server:

pub async fn run(listener: TcpListener, core: Core)

This function should encapsulate things like routing configuration, server logging configuration and so on.

Core can be provided to handlers using Extension Layer mechanism like this:

...
let shared_core = Arc::new(core);
...
let app = Router::new()
    .route("/something", post(post_something))
    ...
    .layer(AddExtensionLayer::new(shared_core));

Which in a handler can be declared in parameter list using extension extractor:

async fn post_something(
    Extension(core): Extension<Arc<Core>>,
    Json(new_something): Json<NewSomething>,
) -> impl IntoResponse {
    core
        .create_something(new_something)
        .await
}

Axum examples contain one on error handling and dependency injection. You can check it here.

Last but not least, now you can mock Core out with a library like mockall, write spawn_app function that would return host and port where the server is run, run some requests against it and do assertions.

The video from Bogdan at Let's Get Rusty channel provides a good start with mockall.

I will be happy to provide more details if you feel something is missing from the answer.

I never thought to provide the core logic as an extension. Clever! Does it scale (in terms of maintainability) as the app grows in complexity? Are there any pitfalls to be aware of? And with the introduction of `State`, would you recommend its use over `Extension`? — verified_tinker, Feb 20 '23 at 17:51
Why did you pass `core` as an `Extension` instead of a `State`? — Patrick, Apr 26 '23 at 08:53

How to design an Axum server with test-friendliness in mind?

1 Answers1