Special handling for HEAD requests

For ever GET route axum automatically installs the same HEAD route.

However we can create special handling for HEAD requests.

It rarely needed.

[package]
name = "head-request"
version = "0.1.0"
edition = "2024"
publish = false


[dependencies]
axum = "0.8.8"
tokio = { version = "1.50.0", features = ["full"] }

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

use axum::response::{Html, IntoResponse, Response};
use axum::{Router, http, routing::get, routing::head};

async fn main_page() -> Response {
    let content = String::from(
        r#"
        <h1>HEAD</h1>
    "#,
    );
    ([("x-my-header", "header of main page")], Html(content)).into_response()
}

async fn just_head() -> Response {
    ([("x-my-header", "header from HEAD of just_head")]).into_response()
}

async fn get_head_handler(method: http::Method) -> Response {
    if method == http::Method::HEAD {
        return ([("x-my-header", "header for HEAD get_head_handler")]).into_response();
    }

    println!("do some heavy computing task in GET");

    (
        [("x-my-header", "header for GET get_head_handler")],
        Html("The content"),
    )
        .into_response()
}

fn create_route() -> Router {
    Router::new()
        .route("/", get(main_page))
        .route("/just-head", head(just_head))
        .route("/get-head", get(get_head_handler))
}

#[tokio::main]
async fn main() {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
        .await
        .unwrap();
    println!("listening on {}", listener.local_addr().unwrap());
    axum::serve(listener, create_route()).await.unwrap();
}

#[cfg(test)]
mod tests;

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

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

    let response = app
        .oneshot(Request::get("/").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");

    assert_eq!(
        response.headers().get("x-my-header").unwrap(),
        "header of main page"
    );

    let bytes = response.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();
    assert!(html.contains("<h1>HEAD</h1>"));
}

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

    let response = app
        .oneshot(Request::head("/").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");

    assert_eq!(
        response.headers().get("x-my-header").unwrap(),
        "header of main page"
    );

    let bytes = response.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();
    assert_eq!(html, "");
}

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

    let response = app
        .oneshot(Request::get("/get-head").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response.headers()["x-my-header"],
        "header for GET get_head_handler"
    );

    let bytes = response.collect().await.unwrap().to_bytes();
    let html = String::from_utf8(bytes.to_vec()).unwrap();
    assert_eq!(html, "The content");
}

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

    let response = app
        .oneshot(Request::head("/get-head").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response.headers()["x-my-header"],
        "header for HEAD get_head_handler"
    );

    let body = response.collect().await.unwrap().to_bytes();
    assert!(body.is_empty());
}

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

    let response = app
        .oneshot(Request::head("/just-head").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response.headers()["x-my-header"],
        "header from HEAD of just_head"
    );

    let body = response.collect().await.unwrap().to_bytes();
    assert!(body.is_empty());
}

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

    let response = app
        .oneshot(Request::get("/just-head").body(Body::empty()).unwrap())
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);

    let body = response.collect().await.unwrap().to_bytes();
    assert!(body.is_empty());
}
}