diff --git a/crates/iota-storage/tests/key_value_tests.rs b/crates/iota-storage/tests/key_value_tests.rs index 8f20904994f..07fa3b01421 100644 --- a/crates/iota-storage/tests/key_value_tests.rs +++ b/crates/iota-storage/tests/key_value_tests.rs @@ -427,42 +427,47 @@ async fn test_get_tx_from_fallback() { #[cfg(msim)] mod simtests { use std::{ + net::SocketAddr, sync::Mutex, time::{Duration, Instant}, }; use axum::{ + body::Body, extract::{Request, State}, + response::Response, routing::get, }; use iota_macros::sim_test; use iota_simulator::configs::constant_latency_ms; use iota_storage::http_key_value_store::*; use rustls::crypto::{CryptoProvider, ring}; - use tokio::net::TcpListener; use tracing::info; use super::*; + async fn svc( + State(state): State>>>>, + request: Request, + ) -> Response { + let path = request.uri().path().to_string(); + let key = path.trim_start_matches('/'); + let value = state.lock().unwrap().get(key).cloned(); + info!("Got request for key: {:?}, value: {:?}", key, value); + match value { + Some(v) => Response::new(Body::from(v)), + None => Response::builder() + .status(hyper::StatusCode::NOT_FOUND) + .body(Body::empty()) + .unwrap(), + } + } + async fn test_server(data: Arc>>>) { let handle = iota_simulator::runtime::Handle::current(); let builder = handle.create_node(); let (startup_sender, mut startup_receiver) = tokio::sync::watch::channel(false); let startup_sender = Arc::new(startup_sender); - let (sender, _) = tokio::sync::broadcast::channel::<()>(1); - - async fn get_data( - data: State>>>>, - req: Request, - ) -> Result, String> { - let path = req.uri().path().to_string(); - let key = path.trim_start_matches('/'); - let value = data.lock().unwrap().get(key).cloned(); - info!("Got request for key: {:?}, value: {:?}", key, value); - value.ok_or_else(|| "no value".to_owned()) - } - - let sender_clone = sender.clone(); let _node = builder .ip("10.10.10.10".parse().unwrap()) .name("server") @@ -470,21 +475,13 @@ mod simtests { info!("Server started"); let data = data.clone(); let startup_sender = startup_sender.clone(); - let mut receiver = sender_clone.subscribe(); async move { - let app = axum::Router::new() - .route("/", get(get_data)) - .with_state(data); - - let addr = TcpListener::bind(("10.10.10.10", 8080)).await.unwrap(); - - tokio::spawn(async move { - axum::serve(addr, app) - .with_graceful_shutdown(async move { - receiver.recv().await.ok(); - }) - .await - .unwrap() + let router = get(svc).with_state(data); + let addr = SocketAddr::from(([10, 10, 10, 10], 8080)); + let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + + tokio::spawn(async { + axum::serve(listener, router).await.unwrap(); }); startup_sender.send(true).ok(); @@ -492,7 +489,6 @@ mod simtests { }) .build(); startup_receiver.changed().await.unwrap(); - sender.send(()).ok(); } #[sim_test(config = "constant_latency_ms(250)")]