Counter with Tower-cookies

The dependencies in Cargo.toml

[package]
name = "counter-with-tower-cookies"
version = "0.1.0"
edition = "2024"

[dependencies]
axum = "0.8.7"
tokio = { version = "1.48.0", features = ["full"] }
tower-cookies = "0.11.0"

[dev-dependencies]
http-body-util = "0.1.3"
tower = { version = "0.5.2", features = ["util"] }

The code

use axum::{Router, response::Html, routing::get};
use std::net::SocketAddr;
use tower_cookies::{Cookie, CookieManagerLayer, Cookies};

const COOKIE_NAME: &str = "visited";

#[tokio::main]
async fn main() {
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
    println!("Listening on http://{}", addr);
    axum::serve(listener, app().into_make_service())
        .await
        .unwrap();
}

fn app() -> Router {
    Router::new()
        .route("/", get(increment))
        .route("/remove", get(remove))
        .layer(CookieManagerLayer::new())
}
async fn increment(cookies: Cookies) -> Html<String> {
    let visited = cookies
        .get(COOKIE_NAME)
        .and_then(|c| c.value().parse().ok())
        .unwrap_or(0);
    cookies.add(Cookie::new(COOKIE_NAME, (visited + 1).to_string()));
    Html(format!(
        r#"You've been here {visited} times before.
    Reload the page to increment the counter or click here to remove the cookie and <a href="/remove">reset the counter</a>."#
    ))
}

async fn remove(cookies: Cookies) -> Html<&'static str> {
    cookies.remove(Cookie::new(COOKIE_NAME, ""));
    Html(r#"Counter has been reset. <a href="/">Go back</a>"#.into())
}

#[cfg(test)]
mod tests;

Tests

#![allow(unused)]
fn main() {
use super::*;
use axum::{body::Body, http::Request, http::StatusCode};
use http_body_util::BodyExt;
use tower::{Service, ServiceExt};

// Actually we don't need the multi-step test here, but it's a good demonstration of how to do it
#[tokio::test]
async fn test_counter_multi_step() {
    let mut svc = app().into_service();

    let response = svc
        .call(Request::builder().uri("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    let cookie_header = response.headers().get("set-cookie").unwrap();
    let cookie_value = cookie_header.to_str().unwrap().to_string();
    assert_eq!(cookie_value, "visited=1");

    let body = response.into_body();
    let bytes = body.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();
    assert!(html.contains(r#"You've been here 0 times before."#));

    let response = svc
        .call(
            Request::builder()
                .uri("/")
                .header("cookie", cookie_value)
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    let cookie_header = response.headers().get("set-cookie").unwrap();
    let cookie_value = cookie_header.to_str().unwrap().to_string();
    assert_eq!(cookie_value, "visited=2");

    let body = response.into_body();
    let bytes = body.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();

    assert!(html.contains(r#"You've been here 1 times before."#));
}

#[tokio::test]
async fn without_cookie() {
    let app = app();

    let response = app
        .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    let cookie_header = response.headers().get("set-cookie").unwrap();
    let cookie_value = cookie_header.to_str().unwrap().to_string();
    assert_eq!(cookie_value, "visited=1");

    let body = response.into_body();
    let bytes = body.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();

    assert!(html.contains(r#"You've been here 0 times before."#));
}

#[tokio::test]
async fn with_cookie() {
    let app = app();

    let response = app
        .oneshot(
            Request::builder()
                .uri("/")
                .header("cookie", "visited=41")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    let cookie_header = response.headers().get("set-cookie").unwrap();
    let cookie_value = cookie_header.to_str().unwrap().to_string();
    assert_eq!(cookie_value, "visited=42");

    let body = response.into_body();
    let bytes = body.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();

    assert!(html.contains(r#"You've been here 41 times before."#));
}

#[tokio::test]
async fn remove_cookie() {
    let app = app();

    let response = app
        .oneshot(
            Request::builder()
                .uri("/remove")
                .header("cookie", "visited=23")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    let cookie_header = response.headers().get("set-cookie").unwrap();
    let cookie_value = cookie_header.to_str().unwrap().to_string();
    assert!(cookie_value.starts_with("visited=;"));

    let body = response.into_body();
    let bytes = body.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();

    assert_eq!(html, r#"Counter has been reset. <a href="/">Go back</a>"#);
}
}