spikard_http/server/
fast_router.rs1use ahash::AHashMap;
9use axum::body::Body;
10use axum::http::Method;
11
12use crate::handler_trait::StaticResponse;
13
14#[derive(Clone)]
22pub struct FastRouter {
23 routes: AHashMap<Method, AHashMap<String, StaticResponse>>,
24}
25
26impl FastRouter {
27 pub fn new() -> Self {
29 Self {
30 routes: AHashMap::new(),
31 }
32 }
33
34 pub fn insert(&mut self, method: Method, path: &str, resp: &StaticResponse) {
36 self.routes
37 .entry(method)
38 .or_default()
39 .insert(path.to_owned(), resp.clone());
40 }
41
42 pub fn has_routes(&self) -> bool {
44 !self.routes.is_empty()
45 }
46
47 pub fn lookup(&self, method: &Method, path: &str) -> Option<axum::response::Response<Body>> {
53 let by_path = self.routes.get(method)?;
54 let resp = by_path.get(path)?;
55 Some(resp.to_response())
56 }
57}
58
59impl Default for FastRouter {
60 fn default() -> Self {
61 Self::new()
62 }
63}
64
65#[cfg(test)]
66mod tests {
67 use super::*;
68 use axum::http::{HeaderValue, StatusCode};
69 use bytes::Bytes;
70 use http_body_util::BodyExt;
71
72 fn make_static_response(status: u16, body: &str) -> StaticResponse {
73 StaticResponse {
74 status,
75 headers: vec![],
76 body: Bytes::from(body.to_owned()),
77 content_type: HeaderValue::from_static("text/plain"),
78 }
79 }
80
81 #[test]
82 fn test_fast_router_miss_returns_none() {
83 let router = FastRouter::new();
84 assert!(router.lookup(&Method::GET, "/health").is_none());
85 }
86
87 #[test]
88 fn test_fast_router_hit_returns_response() {
89 let mut router = FastRouter::new();
90 router.insert(Method::GET, "/health", &make_static_response(200, "OK"));
91
92 let resp = router.lookup(&Method::GET, "/health");
93 assert!(resp.is_some());
94 let resp = resp.unwrap();
95 assert_eq!(resp.status(), StatusCode::OK);
96 }
97
98 #[test]
99 fn test_fast_router_method_mismatch() {
100 let mut router = FastRouter::new();
101 router.insert(Method::GET, "/health", &make_static_response(200, "OK"));
102
103 assert!(router.lookup(&Method::POST, "/health").is_none());
104 }
105
106 #[test]
107 fn test_fast_router_path_mismatch() {
108 let mut router = FastRouter::new();
109 router.insert(Method::GET, "/health", &make_static_response(200, "OK"));
110
111 assert!(router.lookup(&Method::GET, "/ready").is_none());
112 }
113
114 #[test]
115 fn test_fast_router_has_routes() {
116 let mut router = FastRouter::new();
117 assert!(!router.has_routes());
118
119 router.insert(Method::GET, "/health", &make_static_response(200, "OK"));
120 assert!(router.has_routes());
121 }
122
123 #[test]
124 fn test_fast_router_multiple_routes() {
125 let mut router = FastRouter::new();
126 router.insert(Method::GET, "/health", &make_static_response(200, "OK"));
127 router.insert(Method::GET, "/ready", &make_static_response(200, "ready"));
128 router.insert(Method::POST, "/health", &make_static_response(201, "created"));
129
130 assert!(router.lookup(&Method::GET, "/health").is_some());
131 assert!(router.lookup(&Method::GET, "/ready").is_some());
132 assert!(router.lookup(&Method::POST, "/health").is_some());
133 assert!(router.lookup(&Method::DELETE, "/health").is_none());
134 }
135
136 #[test]
137 fn test_fast_router_custom_headers() {
138 use axum::http::header::HeaderName;
139
140 let resp = StaticResponse {
141 status: 200,
142 headers: vec![(HeaderName::from_static("x-custom"), HeaderValue::from_static("value"))],
143 body: Bytes::from_static(b"OK"),
144 content_type: HeaderValue::from_static("application/json"),
145 };
146
147 let mut router = FastRouter::new();
148 router.insert(Method::GET, "/test", &resp);
149
150 let response = router.lookup(&Method::GET, "/test").unwrap();
151 assert_eq!(response.headers().get("x-custom").unwrap(), "value");
152 assert_eq!(response.headers().get("content-type").unwrap(), "application/json");
153 }
154
155 #[tokio::test]
156 async fn test_fast_router_response_body_content() {
157 let mut router = FastRouter::new();
158 router.insert(Method::GET, "/health", &make_static_response(200, "OK"));
159 router.insert(Method::GET, "/ready", &make_static_response(200, "ready"));
160
161 let resp = router.lookup(&Method::GET, "/health").unwrap();
162 let body = resp.into_body().collect().await.unwrap().to_bytes();
163 assert_eq!(body.as_ref(), b"OK");
164
165 let resp = router.lookup(&Method::GET, "/ready").unwrap();
166 let body = resp.into_body().collect().await.unwrap().to_bytes();
167 assert_eq!(body.as_ref(), b"ready");
168 }
169
170 #[tokio::test]
171 async fn test_fast_router_custom_status_code() {
172 let mut router = FastRouter::new();
173 router.insert(Method::POST, "/items", &make_static_response(201, "created"));
174
175 let resp = router.lookup(&Method::POST, "/items").unwrap();
176 assert_eq!(resp.status(), StatusCode::CREATED);
177 let body = resp.into_body().collect().await.unwrap().to_bytes();
178 assert_eq!(body.as_ref(), b"created");
179 }
180
181 #[test]
182 fn test_fast_router_default() {
183 let router = FastRouter::default();
184 assert!(!router.has_routes());
185 }
186}