From 2dc164b4f82aa0436ed0838bf8807ca02f80b3e3 Mon Sep 17 00:00:00 2001
From: Yann Simon <yann.simon@commercetools.com>
Date: Thu, 10 Oct 2024 11:34:01 +0200
Subject: [PATCH] Add test to cover state cloning in layer

---
 axum/src/routing/tests/mod.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs
index c9521e02..af8dca1f 100644
--- a/axum/src/routing/tests/mod.rs
+++ b/axum/src/routing/tests/mod.rs
@@ -3,6 +3,7 @@ use crate::{
     error_handling::HandleErrorLayer,
     extract::{self, DefaultBodyLimit, FromRef, Path, State},
     handler::{Handler, HandlerWithoutStateExt},
+    middleware::{self, Next},
     response::{IntoResponse, Response},
     routing::{
         delete, get, get_service, on, on_service, patch, patch_service,
@@ -924,6 +925,26 @@ async fn state_isnt_cloned_too_much() {
     assert_eq!(state.count(), 3);
 }
 
+#[crate::test]
+async fn state_isnt_cloned_too_much_in_layer() {
+    async fn layer(State(_): State<CountingCloneableState>, req: Request, next: Next) -> Response {
+        next.run(req).await
+    }
+
+    let state = CountingCloneableState::new();
+
+    let app = Router::new().layer(middleware::from_fn_with_state(state.clone(), layer));
+
+    let client = TestClient::new(app);
+
+    // ignore clones made during setup
+    state.setup_done();
+
+    client.get("/").await;
+
+    assert_eq!(state.count(), 4);
+}
+
 #[crate::test]
 async fn logging_rejections() {
     #[derive(Deserialize, Eq, PartialEq, Debug)]