GreptimeTeam · killme2008 · Apr 9, 2024 · Apr 8, 2024 · Apr 9, 2024
@@ -133,6 +133,7 @@ reqwest = { version = "0.11", default-features = false, features = [
     "json",
     "rustls-tls-native-roots",
     "stream",
+    "multipart",
 ] }
 rskafka = "0.5"
 rust_decimal = "1.33"

@@ -51,6 +51,7 @@ futures = "0.3"
 hashbrown = "0.14"
 headers = "0.3"
 hostname = "0.3.1"
+http = "0.2.12"
 http-body = "0.4"
 humantime-serde.workspace = true
 hyper = { version = "0.14", features = ["full"] }
@@ -109,7 +110,6 @@ tikv-jemalloc-ctl = { version = "0.5", features = ["use_std"] }
 
 [dev-dependencies]
 auth = { workspace = true, features = ["testing"] }
-axum-test-helper = "0.3"
 catalog = { workspace = true, features = ["testing"] }
 client.workspace = true
 common-base.workspace = true

@@ -94,6 +94,9 @@ pub mod greptime_result_v1;
 pub mod influxdb_result_v1;
 pub mod table_result;
 
+#[cfg(any(test, feature = "testing"))]
+pub mod test_helpers;
+
 pub const HTTP_API_VERSION: &str = "v1";
 pub const HTTP_API_PREFIX: &str = "/v1/";
 /// Default http body limit (64M).
@@ -824,7 +827,6 @@ mod test {
     use axum::handler::Handler;
     use axum::http::StatusCode;
     use axum::routing::get;
-    use axum_test_helper::TestClient;
     use common_query::Output;
     use common_recordbatch::RecordBatches;
     use datatypes::prelude::*;
@@ -838,6 +840,7 @@ mod test {
 
     use super::*;
     use crate::error::Error;
+    use crate::http::test_helpers::TestClient;
     use crate::query_handler::grpc::GrpcQueryHandler;
     use crate::query_handler::sql::{ServerSqlQueryHandlerAdapter, SqlQueryHandler};
 

@@ -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 convenient 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
+    }
+}
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 use axum::Router;
-use axum_test_helper::TestClient;
 use common_test_util::ports;
+use servers::http::test_helpers::TestClient;
 use servers::http::{HttpOptions, HttpServerBuilder};
 use table::test_util::MemTable;
 

@@ -19,14 +19,14 @@ use api::v1::RowInsertRequests;
 use async_trait::async_trait;
 use auth::tests::{DatabaseAuthInfo, MockUserProvider};
 use axum::{http, Router};
-use axum_test_helper::TestClient;
 use common_query::Output;
 use common_test_util::ports;
 use query::parser::PromQuery;
 use query::plan::LogicalPlan;
 use query::query_engine::DescribeResult;
 use servers::error::{Error, Result};
 use servers::http::header::constants::GREPTIME_DB_HEADER_NAME;
+use servers::http::test_helpers::TestClient;
 use servers::http::{HttpOptions, HttpServerBuilder};
 use servers::influxdb::InfluxdbRequest;
 use servers::query_handler::grpc::GrpcQueryHandler;

@@ -17,13 +17,13 @@ use std::sync::Arc;
 use api::v1::greptime_request::Request;
 use async_trait::async_trait;
 use axum::Router;
-use axum_test_helper::TestClient;
 use common_query::Output;
 use common_test_util::ports;
 use query::parser::PromQuery;
 use query::plan::LogicalPlan;
 use query::query_engine::DescribeResult;
 use servers::error::{self, Result};
+use servers::http::test_helpers::TestClient;
 use servers::http::{HttpOptions, HttpServerBuilder};
 use servers::opentsdb::codec::DataPoint;
 use servers::query_handler::grpc::GrpcQueryHandler;

@@ -21,7 +21,6 @@ use api::v1::greptime_request::Request;
 use api::v1::RowInsertRequests;
 use async_trait::async_trait;
 use axum::Router;
-use axum_test_helper::TestClient;
 use common_query::Output;
 use common_test_util::ports;
 use prost::Message;
@@ -30,6 +29,7 @@ use query::plan::LogicalPlan;
 use query::query_engine::DescribeResult;
 use servers::error::{Error, Result};
 use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF};
+use servers::http::test_helpers::TestClient;
 use servers::http::{HttpOptions, HttpServerBuilder};
 use servers::prom_store;
 use servers::prom_store::{snappy_compress, Metrics};

@@ -16,7 +16,6 @@ arrow-flight.workspace = true
 async-trait = "0.1"
 auth.workspace = true
 axum.workspace = true
-axum-test-helper = "0.3.0"
 catalog.workspace = true
 chrono.workspace = true
 client = { workspace = true, features = ["testing"] }

@@ -17,7 +17,6 @@ use std::collections::BTreeMap;
 use api::prom_store::remote::WriteRequest;
 use auth::user_provider_from_option;
 use axum::http::{HeaderName, StatusCode};
-use axum_test_helper::TestClient;
 use common_error::status_code::StatusCode as ErrorCode;
 use prost::Message;
 use serde_json::json;
@@ -27,6 +26,7 @@ use servers::http::handler::HealthResponse;
 use servers::http::header::GREPTIME_TIMEZONE_HEADER_NAME;
 use servers::http::influxdb_result_v1::{InfluxdbOutput, InfluxdbV1Response};
 use servers::http::prometheus::{PrometheusJsonResponse, PrometheusResponse};
+use servers::http::test_helpers::TestClient;
 use servers::http::GreptimeQueryOutput;
 use servers::prom_store;
 use tests_integration::test_util::{