axum-server: Add unit test for /hello route

[learn-rust.git] / axum-server / src / main.rs

diff --git a/axum-server/src/main.rs b/axum-server/src/main.rs
index 2a2eaaab2f63339bbd80235ecf29337af63503e6..7bb1e94fc0cc41f990ec4646eb22d6fd4c26c366 100644 (file)
--- a/axum-server/src/main.rs
+++ b/axum-server/src/main.rs
@@ -42,10 +42,8 @@ async fn ws_echo(mut socket: ws::WebSocket) {
     }
 }
 
-#[tokio::main]
-async fn main() {
-    tracing_subscriber::fmt::init();
-    let app = Router::new()
+fn app() -> Router {
+    Router::new()
         .route("/hello/:name", get(hello))
         .nest("/dir",
                get_service(tower_http::services::ServeDir::new("../static").precompressed_gzip())
@@ -60,12 +58,52 @@ async fn main() {
                     tower_http::trace::TraceLayer::new_for_http()
                          .make_span_with(tower_http::trace::DefaultMakeSpan::default().include_headers(true)),
                 )
-        );
+        )
+}
+
+#[tokio::main]
+async fn main() {
+    tracing_subscriber::fmt::init();
 
     let addr = std::net::SocketAddr::from(([127, 0, 0, 1], 3030));
     tracing::info!("listening on {}", addr);
     axum::Server::bind(&addr)
-        .serve(app.into_make_service())
+        .serve(app().into_make_service())
         .await
         .unwrap();
 }
+
+#[cfg(test)]
+mod tests {
+    use axum::{
+        http::{Request, StatusCode},
+        response::Response,
+        body::Body
+    };
+    use tower::ServiceExt; // for `oneshot`
+
+    async fn assert_res_ok_body(res: Response, expected_body: &[u8]) {
+        assert_eq!(res.status(), StatusCode::OK);
+        assert_eq!(hyper::body::to_bytes(res.into_body()).await.unwrap(), expected_body);
+    }
+
+    #[tokio::test]
+    async fn test_hello() {
+        // no user-agent
+        let res = super::app()
+            .oneshot(Request::builder().uri("/hello/rust").body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_res_ok_body(res, b"Hello rust").await;
+
+        // with user-agent
+        let res = super::app()
+            .oneshot(Request::builder()
+                        .uri("/hello/rust")
+                        .header("user-agent", "TestBrowser 0.1")
+                        .body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_res_ok_body(res, b"Hello rust from TestBrowser 0.1").await;
+    }
+}