Session: counter with cookie

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

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

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

Code

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

const COOKIE_NAME: &str = "counter";

async fn handle_main_page(cookies: Cookies) -> Html<String> {
    let mut counter: u32 = match cookies.get(COOKIE_NAME) {
        Some(cookie) => cookie.value().parse::<u32>().unwrap(),
        None => 0,
    };
    counter += 1;

    cookies.add(Cookie::new(COOKIE_NAME, counter.to_string()));

    Html(format!(
        r#"<h1>Count {counter}</h1><a href="/delete">delete</a>"#
    ))
}

async fn delete_cookie(cookies: Cookies) -> Html<String> {
    cookies.remove(Cookie::new(COOKIE_NAME, ""));

    Html(format!(r#"<a href="/">home</a>"#))
}

fn create_router() -> Router {
    Router::new()
        .route("/", get(handle_main_page))
        .route("/delete", get(delete_cookie))
        .layer(CookieManagerLayer::new())
}

#[tokio::main]
async fn main() {
    let app = create_router();

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
        .await
        .unwrap();
    println!("listening on http://{}", listener.local_addr().unwrap());
    axum::serve(listener, app).await.unwrap();
}

#[cfg(test)]
mod tests;

Test

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

use super::*;

#[tokio::test]
async fn test_main_page() {
    let response = create_router()
        .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

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

    let headers = response.headers().clone();
    let content_type = headers.get("content-type").unwrap();
    assert_eq!(content_type.to_str().unwrap(), "text/html; charset=utf-8");

    let cookies = headers.get_all("set-cookie");
    let counter_cookie = cookies
        .iter()
        .find(|c| c.to_str().unwrap().contains("counter="))
        .expect("Counter cookie should be set");
    assert_eq!(counter_cookie, "counter=1");

    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, "<h1>Count 1</h1><a href=\"/delete\">delete</a>");

    // new request, now with cookie
    let response = create_router()
        .oneshot(
            Request::builder()
                .uri("/")
                .header("Cookie", counter_cookie)
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

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

    let content_type = response.headers().get("content-type").unwrap();
    assert_eq!(content_type.to_str().unwrap(), "text/html; charset=utf-8");

    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, "<h1>Count 2</h1><a href=\"/delete\">delete</a>");
}
}