Add tests for the mirror-server handling of requests.

This commit is contained in:
Dorian 2024-02-13 16:25:03 -05:00 committed by Dorian Puła
parent 05bffc7996
commit 264cd468e3
2 changed files with 147 additions and 18 deletions

View File

@ -25,3 +25,7 @@ tower = "0.4"
tower-http = { version = "0.5", features = ["trace"] } tower-http = { version = "0.5", features = ["trace"] }
tracing = "0.1" tracing = "0.1"
tracing-subscriber = "0.3" tracing-subscriber = "0.3"
[dev-dependencies]
axum-test = "14.0"
rstest = "0.18"

View File

@ -1,14 +1,14 @@
use axum::{ use axum::{
extract::{Host, OriginalUri, Json}, extract::{Host, Json, OriginalUri},
http::{header::HeaderMap, Method}, http::{header::HeaderMap, Method},
Router, Router,
}; };
use clap::Parser; use clap::Parser;
use serde::Serialize; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use std::collections::BTreeMap; use std::collections::BTreeMap;
use tower_http::trace; use tower_http::trace;
use tracing::{Level, info, warn}; use tracing::{info, warn, Level};
#[derive(Parser)] #[derive(Parser)]
#[command(author, version, about, long_about = None)] #[command(author, version, about, long_about = None)]
@ -22,7 +22,7 @@ struct CliArgs {
ips: String, ips: String,
} }
#[derive(Serialize)] #[derive(Serialize, Deserialize, Debug, PartialEq)]
struct EchoResponse { struct EchoResponse {
method: String, method: String,
path: String, path: String,
@ -37,7 +37,7 @@ async fn echo_request(
original_uri: OriginalUri, original_uri: OriginalUri,
host: Host, host: Host,
header_map: HeaderMap, header_map: HeaderMap,
body: Option<Json<Value>> body: Option<Json<Value>>,
) -> Json<EchoResponse> { ) -> Json<EchoResponse> {
let method = method.to_string(); let method = method.to_string();
let host = host.0; let host = host.0;
@ -47,25 +47,34 @@ async fn echo_request(
.map(|(name, value)| (name.to_string(), value.to_str().unwrap_or("").to_string())) .map(|(name, value)| (name.to_string(), value.to_str().unwrap_or("").to_string()))
.collect(); .collect();
let json_body = match body { let body = match body {
None => { None => {
warn!("Received a non-JSON body."); warn!("Received a non-JSON body.");
None None
}, }
Some(Json(value)) => { Some(Json(value)) => {
info!("JSON request: {}", value.to_string()); info!("JSON request: {}", value.to_string());
Some(value)}, Some(value)
}
}; };
let response = EchoResponse { let response = EchoResponse {
method, method,
host, host,
path, path,
headers, headers,
body: json_body, body,
}; };
Json(response) Json(response)
} }
fn app() -> Router {
Router::new().fallback(echo_request).layer(
trace::TraceLayer::new_for_http()
.make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
.on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
)
}
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
let cli_args = CliArgs::parse(); let cli_args = CliArgs::parse();
@ -73,20 +82,136 @@ async fn main() {
let listen_on = format!("{ip}:{port}", ip = listen_on.0, port = listen_on.1); let listen_on = format!("{ip}:{port}", ip = listen_on.0, port = listen_on.1);
// From https://stackoverflow.com/questions/75009289/how-to-enable-logging-tracing-with-axum // From https://stackoverflow.com/questions/75009289/how-to-enable-logging-tracing-with-axum
tracing_subscriber::fmt() tracing_subscriber::fmt().with_max_level(Level::INFO).init();
.with_max_level(Level::INFO)
.init();
let app = Router::new().fallback(echo_request).layer(
trace::TraceLayer::new_for_http()
.make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
.on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
);
info!("Starting the mirror-server to listen to {}", listen_on); info!("Starting the mirror-server to listen to {}", listen_on);
let listener = tokio::net::TcpListener::bind(&listen_on) let listener = tokio::net::TcpListener::bind(&listen_on)
.await .await
.unwrap_or_else(|_| panic!("Failed to binding to {}", listen_on)); .unwrap_or_else(|_| panic!("Failed to binding to {}", listen_on));
axum::serve(listener, app) axum::serve(listener, app())
.await .await
.expect("Server should start"); .expect("Server should start");
} }
#[cfg(test)]
mod tests {
use super::*;
use axum::http::{HeaderName, HeaderValue, Method, StatusCode};
use axum_test::TestServer;
use rstest::rstest;
#[tokio::test]
async fn handles_simple_get_request() {
let expected_headers: BTreeMap<String, String> = BTreeMap::new();
let expected_response = EchoResponse {
method: "GET".to_string(),
path: "/".to_string(),
host: "localhost".to_string(),
body: None,
headers: expected_headers,
};
let server = TestServer::new(app()).unwrap();
let response = server.get("/").await;
response.assert_status(StatusCode::OK);
response.assert_json::<EchoResponse>(&expected_response);
}
#[rstest]
#[case::single_level("/test")]
#[case::multiple_level("/test/multiple/levels/")]
#[case::slash_ending("/test/")]
#[tokio::test]
async fn handles_different_urls(#[case] url: String) {
let expected_headers: BTreeMap<String, String> = BTreeMap::new();
let expected_response = EchoResponse {
method: "GET".to_string(),
path: url.clone(),
host: "localhost".to_string(),
body: None,
headers: expected_headers,
};
let server = TestServer::new(app()).unwrap();
let response = server.get(&url).await;
response.assert_status(StatusCode::OK);
response.assert_json::<EchoResponse>(&expected_response);
}
#[rstest]
#[case::post(Method::POST)]
#[case::put(Method::PUT)]
#[case::patch(Method::PATCH)]
#[case::delete(Method::DELETE)]
#[tokio::test]
async fn handles_different_http_methods(#[case] http_method: Method) {
let expected_headers: BTreeMap<String, String> = BTreeMap::new();
let expected_response = EchoResponse {
method: http_method.to_string(),
path: "/testing".to_string(),
host: "localhost".to_string(),
body: None,
headers: expected_headers,
};
let server = TestServer::new(app()).unwrap();
let response = server.method(http_method, "/testing").await;
response.assert_status(StatusCode::OK);
response.assert_json::<EchoResponse>(&expected_response);
}
#[tokio::test]
async fn handle_non_json_request() {
let mut expected_headers: BTreeMap<String, String> = BTreeMap::new();
expected_headers.insert("content-type".to_string(), "text/plain".to_string());
let expected_response = EchoResponse {
method: "POST".to_string(),
path: "/".to_string(),
host: "localhost".to_string(),
body: None,
headers: expected_headers,
};
let server = TestServer::new(app()).unwrap();
let response = server.post("/").text("hello world").await;
response.assert_status(StatusCode::OK);
response.assert_json::<EchoResponse>(&expected_response);
}
#[tokio::test]
async fn handle_json_request() {
let test_json = serde_json::json!({"hello": "world"});
let mut expected_headers: BTreeMap<String, String> = BTreeMap::new();
expected_headers.insert("content-type".to_string(), "application/json".to_string());
let expected_response = EchoResponse {
method: "POST".to_string(),
path: "/".to_string(),
host: "localhost".to_string(),
body: Some(test_json.clone()),
headers: expected_headers,
};
let server = TestServer::new(app()).unwrap();
let response = server.post("/").json(&test_json).await;
response.assert_status(StatusCode::OK);
response.assert_json::<EchoResponse>(&expected_response);
}
#[tokio::test]
async fn handle_extra_headers() {
let mut expected_headers: BTreeMap<String, String> = BTreeMap::new();
expected_headers.insert("x-test-message".to_string(), "Howdy!".to_string());
let expected_response = EchoResponse {
method: "GET".to_string(),
path: "/".to_string(),
host: "localhost".to_string(),
body: None,
headers: expected_headers,
};
let server = TestServer::new(app()).unwrap();
let response = server
.get("/")
.add_header(
HeaderName::from_static("x-test-message"),
HeaderValue::from_static("Howdy!"),
)
.await;
response.assert_status(StatusCode::OK);
response.assert_json::<EchoResponse>(&expected_response);
}
}