servers/http/
test_helpers.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// This file is copied from https://github.com/tokio-rs/axum/blob/axum-v0.6.20/axum/src/test_helpers/test_client.rs
16
17//! Axum Test Client
18//!
19//! ```rust
20//! use axum::Router;
21//! use axum::http::StatusCode;
22//! use axum::routing::get;
23//! use crate::servers::http::test_helpers::TestClient;
24//!
25//! let async_block = async {
26//!     // you can replace this Router with your own app
27//!     let app = Router::new().route("/", get(|| async {}));
28//!
29//!     // initiate the TestClient with the previous declared Router
30//!     let client = TestClient::new(app).await;
31//!
32//!     let res = client.get("/").await;
33//!     assert_eq!(res.status(), StatusCode::OK);
34//! };
35//!
36//! // Create a runtime for executing the async block. This runtime is local
37//! // to the main function and does not require any global setup.
38//! let runtime = tokio::runtime::Builder::new_current_thread()
39//!     .enable_all()
40//!     .build()
41//!     .unwrap();
42//!
43//! // Use the local runtime to block on the async block.
44//! runtime.block_on(async_block);
45//! ```
46
47use std::convert::TryFrom;
48use std::net::SocketAddr;
49
50use axum::Router;
51use bytes::Bytes;
52use common_telemetry::info;
53use http::header::{HeaderName, HeaderValue};
54use http::{Method, StatusCode};
55use tokio::net::TcpListener;
56
57/// Test client to Axum servers.
58pub struct TestClient {
59    client: reqwest::Client,
60    addr: SocketAddr,
61}
62
63impl TestClient {
64    /// Create a new test client.
65    pub async fn new(svc: Router) -> Self {
66        let listener = TcpListener::bind("127.0.0.1:0")
67            .await
68            .expect("Could not bind ephemeral socket");
69        let addr = listener.local_addr().unwrap();
70        info!("Listening on {}", addr);
71
72        tokio::spawn(async move {
73            axum::serve(listener, svc).await.expect("server error");
74        });
75
76        let client = reqwest::Client::builder()
77            .redirect(reqwest::redirect::Policy::none())
78            .build()
79            .unwrap();
80
81        TestClient { client, addr }
82    }
83
84    /// Returns the base URL (http://ip:port) for this TestClient
85    ///
86    /// this is useful when trying to check if Location headers in responses
87    /// are generated correctly as Location contains an absolute URL
88    pub fn base_url(&self) -> String {
89        format!("http://{}", self.addr)
90    }
91
92    /// Create a GET request.
93    pub fn get(&self, url: &str) -> RequestBuilder {
94        common_telemetry::info!("GET {} {}", self.addr, url);
95
96        RequestBuilder {
97            builder: self.client.get(format!("http://{}{}", self.addr, url)),
98        }
99    }
100
101    /// Create a HEAD request.
102    pub fn head(&self, url: &str) -> RequestBuilder {
103        common_telemetry::info!("HEAD {} {}", self.addr, url);
104
105        RequestBuilder {
106            builder: self.client.head(format!("http://{}{}", self.addr, url)),
107        }
108    }
109
110    /// Create a POST request.
111    pub fn post(&self, url: &str) -> RequestBuilder {
112        common_telemetry::info!("POST {} {}", self.addr, url);
113
114        RequestBuilder {
115            builder: self.client.post(format!("http://{}{}", self.addr, url)),
116        }
117    }
118
119    /// Create a PUT request.
120    pub fn put(&self, url: &str) -> RequestBuilder {
121        common_telemetry::info!("PUT {} {}", self.addr, url);
122
123        RequestBuilder {
124            builder: self.client.put(format!("http://{}{}", self.addr, url)),
125        }
126    }
127
128    /// Create a PATCH request.
129    pub fn patch(&self, url: &str) -> RequestBuilder {
130        common_telemetry::info!("PATCH {} {}", self.addr, url);
131
132        RequestBuilder {
133            builder: self.client.patch(format!("http://{}{}", self.addr, url)),
134        }
135    }
136
137    /// Create a DELETE request.
138    pub fn delete(&self, url: &str) -> RequestBuilder {
139        common_telemetry::info!("DELETE {} {}", self.addr, url);
140
141        RequestBuilder {
142            builder: self.client.delete(format!("http://{}{}", self.addr, url)),
143        }
144    }
145
146    /// Options preflight request
147    pub fn options(&self, url: &str) -> RequestBuilder {
148        common_telemetry::info!("OPTIONS {} {}", self.addr, url);
149
150        RequestBuilder {
151            builder: self
152                .client
153                .request(Method::OPTIONS, format!("http://{}{}", self.addr, url)),
154        }
155    }
156}
157
158/// Builder for test requests.
159pub struct RequestBuilder {
160    builder: reqwest::RequestBuilder,
161}
162
163impl RequestBuilder {
164    pub async fn send(self) -> TestResponse {
165        TestResponse {
166            response: self.builder.send().await.unwrap(),
167        }
168    }
169
170    /// Set the request body.
171    pub fn body(mut self, body: impl Into<reqwest::Body>) -> Self {
172        self.builder = self.builder.body(body);
173        self
174    }
175
176    /// Set the request forms.
177    pub fn form<T: serde::Serialize + ?Sized>(mut self, form: &T) -> Self {
178        self.builder = self.builder.form(&form);
179        self
180    }
181
182    /// Set the request JSON body.
183    pub fn json<T>(mut self, json: &T) -> Self
184    where
185        T: serde::Serialize,
186    {
187        self.builder = self.builder.json(json);
188        self
189    }
190
191    /// Set a request header.
192    pub fn header<K, V>(mut self, key: K, value: V) -> Self
193    where
194        HeaderName: TryFrom<K>,
195        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
196        HeaderValue: TryFrom<V>,
197        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
198    {
199        self.builder = self.builder.header(key, value);
200
201        self
202    }
203
204    /// Set a request multipart form.
205    pub fn multipart(mut self, form: reqwest::multipart::Form) -> Self {
206        self.builder = self.builder.multipart(form);
207        self
208    }
209}
210
211/// A wrapper around [`reqwest::Response`] that provides common methods with internal `unwrap()`s.
212///
213/// This is convenient for tests where panics are what you want. For access to
214/// non-panicking versions or the complete `Response` API use `into_inner()` or
215/// `as_ref()`.
216#[derive(Debug)]
217pub struct TestResponse {
218    response: reqwest::Response,
219}
220
221impl TestResponse {
222    /// Get the response body as text.
223    pub async fn text(self) -> String {
224        self.response.text().await.unwrap()
225    }
226
227    /// Get the response body as bytes.
228    pub async fn bytes(self) -> Bytes {
229        self.response.bytes().await.unwrap()
230    }
231
232    /// Get the response body as JSON.
233    pub async fn json<T>(self) -> T
234    where
235        T: serde::de::DeserializeOwned,
236    {
237        self.response.json().await.unwrap()
238    }
239
240    /// Get the response status.
241    pub fn status(&self) -> StatusCode {
242        StatusCode::from_u16(self.response.status().as_u16()).unwrap()
243    }
244
245    /// Get the response headers.
246    pub fn headers(&self) -> http::HeaderMap {
247        self.response.headers().clone()
248    }
249
250    /// Get the response in chunks.
251    pub async fn chunk(&mut self) -> Option<Bytes> {
252        self.response.chunk().await.unwrap()
253    }
254
255    /// Get the response in chunks as text.
256    pub async fn chunk_text(&mut self) -> Option<String> {
257        let chunk = self.chunk().await?;
258        Some(String::from_utf8(chunk.to_vec()).unwrap())
259    }
260
261    /// Get the inner [`reqwest::Response`] for less convenient but more complete access.
262    pub fn into_inner(self) -> reqwest::Response {
263        self.response
264    }
265}
266
267impl AsRef<reqwest::Response> for TestResponse {
268    fn as_ref(&self) -> &reqwest::Response {
269        &self.response
270    }
271}
servers/http/test_helpers.rs

servers/http/
test_helpers.rs
