mkv/tests/integration.rs

use reqwest::StatusCode;
use std::sync::atomic::{AtomicU32, Ordering};

static TEST_COUNTER: AtomicU32 = AtomicU32::new(0);

async fn start_server() -> String {
    let id = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
    let db_path = format!("/tmp/mkv-test/index-{id}.db");

    let _ = std::fs::remove_file(&db_path);
    let _ = std::fs::remove_file(format!("{db_path}-wal"));
    let _ = std::fs::remove_file(format!("{db_path}-shm"));

    let args = mkv::Args {
        db_path,
        volumes: vec![
            "http://localhost:3101".into(),
            "http://localhost:3102".into(),
            "http://localhost:3103".into(),
        ],
        replicas: 2,
    };

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let app = mkv::build_app(&args);

    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    format!("http://127.0.0.1:{port}")
}

fn client() -> reqwest::Client {
    reqwest::Client::builder()
        .redirect(reqwest::redirect::Policy::none())
        .build()
        .unwrap()
}

#[tokio::test]
async fn test_put_and_head() {
    let base = start_server().await;
    let c = client();

    let resp = c.put(format!("{base}/hello")).body("world").send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::CREATED);

    let resp = c.head(format!("{base}/hello")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);
    assert_eq!(resp.headers().get("content-length").unwrap().to_str().unwrap(), "5");
}

#[tokio::test]
async fn test_put_and_get_redirect() {
    let base = start_server().await;
    let c = client();

    let resp = c.put(format!("{base}/redirect-test")).body("some data").send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::CREATED);

    let resp = c.get(format!("{base}/redirect-test")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::FOUND);

    let location = resp.headers().get("location").unwrap().to_str().unwrap();
    assert!(location.starts_with("http://localhost:310"), "got: {location}");

    let blob_resp = reqwest::get(location).await.unwrap();
    assert_eq!(blob_resp.status(), StatusCode::OK);
    assert_eq!(blob_resp.text().await.unwrap(), "some data");
}

#[tokio::test]
async fn test_get_nonexistent_returns_404() {
    let base = start_server().await;
    let c = client();
    let resp = c.get(format!("{base}/does-not-exist")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::NOT_FOUND);
}

#[tokio::test]
async fn test_put_get_delete_get() {
    let base = start_server().await;
    let c = client();

    let resp = c.put(format!("{base}/delete-me")).body("temporary").send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::CREATED);

    let resp = c.get(format!("{base}/delete-me")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::FOUND);

    let resp = c.delete(format!("{base}/delete-me")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::NO_CONTENT);

    let resp = c.get(format!("{base}/delete-me")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::NOT_FOUND);
}

#[tokio::test]
async fn test_delete_nonexistent_returns_404() {
    let base = start_server().await;
    let c = client();
    let resp = c.delete(format!("{base}/never-existed")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::NOT_FOUND);
}

#[tokio::test]
async fn test_list_keys() {
    let base = start_server().await;
    let c = client();

    for name in ["docs/a", "docs/b", "docs/c", "other/x"] {
        c.put(format!("{base}/{name}")).body("data").send().await.unwrap();
    }

    let resp = c.get(format!("{base}/")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);
    let body = resp.text().await.unwrap();
    assert!(body.contains("docs/a"));
    assert!(body.contains("other/x"));

    let resp = c.get(format!("{base}/?prefix=docs/")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);
    let body = resp.text().await.unwrap();
    let lines: Vec<&str> = body.lines().collect();
    assert_eq!(lines.len(), 3);
    assert!(!body.contains("other/x"));
}

#[tokio::test]
async fn test_put_overwrite() {
    let base = start_server().await;
    let c = client();

    c.put(format!("{base}/overwrite")).body("version1").send().await.unwrap();

    let resp = c.put(format!("{base}/overwrite")).body("version2").send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::CREATED);

    let resp = c.head(format!("{base}/overwrite")).send().await.unwrap();
    assert_eq!(resp.headers().get("content-length").unwrap().to_str().unwrap(), "8");

    let resp = c.get(format!("{base}/overwrite")).send().await.unwrap();
    let location = resp.headers().get("location").unwrap().to_str().unwrap();
    let body = reqwest::get(location).await.unwrap().text().await.unwrap();
    assert_eq!(body, "version2");
}

#[tokio::test]
async fn test_replication_writes_to_multiple_volumes() {
    let base = start_server().await;
    let c = client();

    c.put(format!("{base}/replicated")).body("replica-data").send().await.unwrap();

    let resp = c.head(format!("{base}/replicated")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let resp = c.get(format!("{base}/replicated")).send().await.unwrap();
    assert_eq!(resp.status(), StatusCode::FOUND);
    let location = resp.headers().get("location").unwrap().to_str().unwrap();
    let body = reqwest::get(location).await.unwrap().text().await.unwrap();
    assert_eq!(body, "replica-data");
}