Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor: bundle the lightweight axum test client
Signed-off-by: tison <wander4096@gmail.com>
- Loading branch information
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,204 @@ | ||
| // Copyright 2023 Greptime Team | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| use std::convert::TryFrom; | ||
| use std::net::{SocketAddr, TcpListener}; | ||
|
|
||
| use axum::body::HttpBody; | ||
| use axum::BoxError; | ||
| use bytes::Bytes; | ||
| use common_telemetry::info; | ||
| use http::header::{HeaderName, HeaderValue}; | ||
| use http::{Request, StatusCode}; | ||
| use hyper::service::Service; | ||
| use hyper::{Body, Server}; | ||
| use tower::make::Shared; | ||
|
|
||
| pub struct TestClient { | ||
| client: reqwest::Client, | ||
| addr: SocketAddr, | ||
| } | ||
|
|
||
| impl TestClient { | ||
| pub fn new<S, ResBody>(svc: S) -> Self | ||
| where | ||
| S: Service<Request<Body>, Response = http::Response<ResBody>> + Clone + Send + 'static, | ||
| ResBody: HttpBody + Send + 'static, | ||
| ResBody::Data: Send, | ||
| ResBody::Error: Into<BoxError>, | ||
| S::Future: Send, | ||
| S::Error: Into<BoxError>, | ||
| { | ||
| let listener = TcpListener::bind("127.0.0.1:0").expect("Could not bind ephemeral socket"); | ||
| let addr = listener.local_addr().unwrap(); | ||
| info!("Listening on {}", addr); | ||
|
|
||
| tokio::spawn(async move { | ||
| let server = Server::from_tcp(listener).unwrap().serve(Shared::new(svc)); | ||
| server.await.expect("server error"); | ||
| }); | ||
|
|
||
| let client = reqwest::Client::builder() | ||
| .redirect(reqwest::redirect::Policy::none()) | ||
| .build() | ||
| .unwrap(); | ||
|
|
||
| TestClient { client, addr } | ||
| } | ||
|
|
||
| /// returns the base URL (http://ip:port) for this TestClient | ||
| /// | ||
| /// this is useful when trying to check if Location headers in responses | ||
| /// are generated correctly as Location contains an absolute URL | ||
| pub fn base_url(&self) -> String { | ||
| format!("http://{}", self.addr) | ||
| } | ||
|
|
||
| pub fn get(&self, url: &str) -> RequestBuilder { | ||
| RequestBuilder { | ||
| builder: self.client.get(format!("http://{}{}", self.addr, url)), | ||
| } | ||
| } | ||
|
|
||
| pub fn head(&self, url: &str) -> RequestBuilder { | ||
| RequestBuilder { | ||
| builder: self.client.head(format!("http://{}{}", self.addr, url)), | ||
| } | ||
| } | ||
|
|
||
| pub fn post(&self, url: &str) -> RequestBuilder { | ||
| RequestBuilder { | ||
| builder: self.client.post(format!("http://{}{}", self.addr, url)), | ||
| } | ||
| } | ||
|
|
||
| pub fn put(&self, url: &str) -> RequestBuilder { | ||
| RequestBuilder { | ||
| builder: self.client.put(format!("http://{}{}", self.addr, url)), | ||
| } | ||
| } | ||
|
|
||
| pub fn patch(&self, url: &str) -> RequestBuilder { | ||
| RequestBuilder { | ||
| builder: self.client.patch(format!("http://{}{}", self.addr, url)), | ||
| } | ||
| } | ||
|
|
||
| pub fn delete(&self, url: &str) -> RequestBuilder { | ||
| RequestBuilder { | ||
| builder: self.client.delete(format!("http://{}{}", self.addr, url)), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pub struct RequestBuilder { | ||
| builder: reqwest::RequestBuilder, | ||
| } | ||
|
|
||
| impl RequestBuilder { | ||
| pub async fn send(self) -> TestResponse { | ||
| TestResponse { | ||
| response: self.builder.send().await.unwrap(), | ||
| } | ||
| } | ||
|
|
||
| pub fn body(mut self, body: impl Into<reqwest::Body>) -> Self { | ||
| self.builder = self.builder.body(body); | ||
| self | ||
| } | ||
|
|
||
| pub fn form<T: serde::Serialize + ?Sized>(mut self, form: &T) -> Self { | ||
| self.builder = self.builder.form(&form); | ||
| self | ||
| } | ||
|
|
||
| pub fn json<T>(mut self, json: &T) -> Self | ||
| where | ||
| T: serde::Serialize, | ||
| { | ||
| self.builder = self.builder.json(json); | ||
| self | ||
| } | ||
|
|
||
| pub fn header<K, V>(mut self, key: K, value: V) -> Self | ||
| where | ||
| HeaderName: TryFrom<K>, | ||
| <HeaderName as TryFrom<K>>::Error: Into<http::Error>, | ||
| HeaderValue: TryFrom<V>, | ||
| <HeaderValue as TryFrom<V>>::Error: Into<http::Error>, | ||
| { | ||
| self.builder = self.builder.header(key, value); | ||
| self | ||
| } | ||
|
|
||
| pub fn multipart(mut self, form: reqwest::multipart::Form) -> Self { | ||
| self.builder = self.builder.multipart(form); | ||
| self | ||
| } | ||
| } | ||
|
|
||
| /// A wrapper around [`reqwest::Response`] that provides common methods with internal `unwrap()`s. | ||
| /// | ||
| /// This is conventient for tests where panics are what you want. For access to | ||
| /// non-panicking versions or the complete `Response` API use `into_inner()` or | ||
| /// `as_ref()`. | ||
| pub struct TestResponse { | ||
| response: reqwest::Response, | ||
| } | ||
|
|
||
| impl TestResponse { | ||
| pub async fn text(self) -> String { | ||
| self.response.text().await.unwrap() | ||
| } | ||
|
|
||
| #[allow(dead_code)] | ||
| pub async fn bytes(self) -> Bytes { | ||
| self.response.bytes().await.unwrap() | ||
| } | ||
|
|
||
| pub async fn json<T>(self) -> T | ||
| where | ||
| T: serde::de::DeserializeOwned, | ||
| { | ||
| self.response.json().await.unwrap() | ||
| } | ||
|
|
||
| pub fn status(&self) -> StatusCode { | ||
| self.response.status() | ||
| } | ||
|
|
||
| pub fn headers(&self) -> &http::HeaderMap { | ||
| self.response.headers() | ||
| } | ||
|
|
||
| pub async fn chunk(&mut self) -> Option<Bytes> { | ||
| self.response.chunk().await.unwrap() | ||
| } | ||
|
|
||
| pub async fn chunk_text(&mut self) -> Option<String> { | ||
| let chunk = self.chunk().await?; | ||
| Some(String::from_utf8(chunk.to_vec()).unwrap()) | ||
| } | ||
|
|
||
| /// Get the inner [`reqwest::Response`] for less convenient but more complete access. | ||
| pub fn into_inner(self) -> reqwest::Response { | ||
| self.response | ||
| } | ||
| } | ||
|
|
||
| impl AsRef<reqwest::Response> for TestResponse { | ||
| fn as_ref(&self) -> &reqwest::Response { | ||
| &self.response | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters