Testing Tokio Services

use std::time::Duration;
use tokio::time::{self, Instant};

async fn retry_with_backoff<F, Fut>(mut op: F, attempts: u32) -> Result<u32, &'static str>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<u32, &'static str>>,
{
    let mut delay = Duration::from_millis(100);
    for _ in 0..attempts {
        match op().await {
            Ok(v) => return Ok(v),
            Err(_) => {
                time::sleep(delay).await;
                delay *= 2;
            }
        }
    }
    Err("gave up")
}

#[tokio::test(start_paused = true)]
async fn backoff_waits_100_200_400ms() {
    let start = Instant::now();
    let mut calls = 0;
    let res = retry_with_backoff(|| { calls += 1; async { Err("down") } }, 3).await;
    assert_eq!(res, Err("gave up"));
    assert_eq!(start.elapsed(), Duration::from_millis(700)); // exact: virtual time
}

#[tokio::test(start_paused = true)]
async fn timeout_fires_without_waiting() {
    let slow = time::sleep(Duration::from_secs(30));
    assert!(time::timeout(Duration::from_secs(5), slow).await.is_err());
}

use axum::{body::{to_bytes, Body}, extract::Path, http::{Request, StatusCode}, routing::get, Json, Router};
use tower::ServiceExt; // oneshot

async fn get_item(Path(id): Path<u64>) -> Result<Json<serde_json::Value>, StatusCode> {
    if id == 0 { Err(StatusCode::NOT_FOUND) } else { Ok(Json(serde_json::json!({ "id": id }))) }
}

fn app() -> Router {
    Router::new().route("/items/{id}", get(get_item))
}

#[tokio::test]
async fn returns_item_json() {
    let res = app().oneshot(Request::get("/items/7").body(Body::empty()).unwrap()).await.unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let body = to_bytes(res.into_body(), 64 * 1024).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    assert_eq!(json["id"], 7);
}

#[tokio::test]
async fn zero_is_not_found() {
    let res = app().oneshot(Request::get("/items/0").body(Body::empty()).unwrap()).await.unwrap();
    assert_eq!(res.status(), StatusCode::NOT_FOUND);
}

use std::{collections::HashMap, sync::{Arc, Mutex}};
use axum::{extract::{Path, State}, http::StatusCode, routing::get, Router};

pub trait Prices: Send + Sync {
    fn price(&self, sku: &str) -> Option<u64>;
}

#[derive(Clone)]
pub struct AppState { prices: Arc<dyn Prices> }

async fn price(State(s): State<AppState>, Path(sku): Path<String>) -> Result<String, StatusCode> {
    s.prices.price(&sku).map(|p| p.to_string()).ok_or(StatusCode::NOT_FOUND)
}

pub fn router(state: AppState) -> Router {
    Router::new().route("/price/{sku}", get(price)).with_state(state)
}

#[derive(Default)]
struct FakePrices(Mutex<HashMap<String, u64>>);

impl Prices for FakePrices {
    fn price(&self, sku: &str) -> Option<u64> {
        self.0.lock().unwrap().get(sku).copied()
    }
}

#[tokio::test]
async fn uses_the_fake() {
    use axum::{body::Body, http::Request};
    use tower::ServiceExt;
    let fake = FakePrices::default();
    fake.0.lock().unwrap().insert("apple".into(), 42);
    let app = router(AppState { prices: Arc::new(fake) });
    let res = app.oneshot(Request::get("/price/apple").body(Body::empty()).unwrap()).await.unwrap();
    assert_eq!(res.status(), StatusCode::OK);
}

use std::time::Duration;
use tokio_util::{sync::CancellationToken, task::TaskTracker};

fn spawn_worker(tracker: &TaskTracker, token: CancellationToken) {
    tracker.spawn(async move {
        let mut tick = tokio::time::interval(Duration::from_secs(1));
        loop {
            tokio::select! {
                _ = token.cancelled() => break,
                _ = tick.tick() => {}
            }
        }
    });
}

#[tokio::test(start_paused = true)]
async fn workers_stop_on_cancel() {
    let (tracker, token) = (TaskTracker::new(), CancellationToken::new());
    for _ in 0..3 {
        spawn_worker(&tracker, token.child_token());
    }
    token.cancel();
    tracker.close();
    tokio::time::timeout(Duration::from_secs(1), tracker.wait()).await.expect("workers stopped");
    assert!(tracker.is_empty());
}

use tokio::sync::mpsc::{self, error::TrySendError};

#[tokio::test]
async fn bounded_channel_pushes_back() {
    let (tx, mut rx) = mpsc::channel::<u32>(2);
    tx.try_send(1).unwrap();
    tx.try_send(2).unwrap();
    assert!(matches!(tx.try_send(3), Err(TrySendError::Full(3)))); // producer must wait or drop
    assert_eq!(rx.recv().await, Some(1));
    tx.try_send(3).unwrap(); // room again
}

fn assert_send<T: Send>(_: &T) {}

async fn handler_logic(state: std::sync::Arc<std::sync::Mutex<Vec<u32>>>) {
    {
        let mut g = state.lock().unwrap();
        g.push(1);
    } // guard dropped before the await
    tokio::task::yield_now().await;
}

#[test]
fn handler_future_is_send() {
    let fut = handler_logic(Default::default());
    assert_send(&fut); // fails to compile if a !Send value lives across an .await
}

#[tokio::test]
async fn serves_over_tcp() {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let token = tokio_util::sync::CancellationToken::new();
    let server = tokio::spawn({
        let token = token.clone();
        async move {
            axum::serve(listener, app()).with_graceful_shutdown(token.cancelled_owned()).await.unwrap();
        }
    });
    let body = reqwest::get(format!("http://{addr}/healthz")).await.unwrap().text().await.unwrap();
    assert_eq!(body, "ok");
    token.cancel();
    server.await.unwrap(); // graceful shutdown completed
}