cloudlysis/gateway/tests/ha_local.rs

use std::sync::Arc;

use tower::util::ServiceExt;

#[tokio::test]
async fn t9_2_and_t9_3_ready_is_healthy_on_both_replicas_and_survives_one_replica_down() {
    let (app1, app2, _state) = build_two_replicas().await;

    let r1 = app1
        .clone()
        .oneshot(
            axum::http::Request::builder()
                .method("GET")
                .uri("/ready")
                .body(axum::body::Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(r1.status(), axum::http::StatusCode::OK);

    let r2 = app2
        .clone()
        .oneshot(
            axum::http::Request::builder()
                .method("GET")
                .uri("/ready")
                .body(axum::body::Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(r2.status(), axum::http::StatusCode::OK);

    drop(app1);

    let r2 = app2
        .oneshot(
            axum::http::Request::builder()
                .method("GET")
                .uri("/ready")
                .body(axum::body::Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(r2.status(), axum::http::StatusCode::OK);
}

#[tokio::test]
async fn t9_4_refresh_works_across_replicas_without_sticky_sessions() {
    let (app1, app2, state) = build_two_replicas().await;

    let signup = app1
        .clone()
        .oneshot(
            axum::http::Request::builder()
                .method("POST")
                .uri("/v1/auth/signup")
                .header("content-type", "application/json")
                .body(axum::body::Body::from(
                    r#"{"email":"ha@b.com","password":"password123"}"#,
                ))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(signup.status(), axum::http::StatusCode::OK);
    let body = axum::body::to_bytes(signup.into_body(), usize::MAX)
        .await
        .unwrap();
    let created: gateway::authn::AuthResponse = serde_json::from_slice(&body).unwrap();

    let refresh_req = serde_json::to_vec(&gateway::authn::RefreshRequest {
        session_id: created.session_id.clone(),
        refresh_token: created.refresh_token.clone(),
    })
    .unwrap();

    let refresh = app2
        .clone()
        .oneshot(
            axum::http::Request::builder()
                .method("POST")
                .uri("/v1/auth/refresh")
                .header("content-type", "application/json")
                .body(axum::body::Body::from(refresh_req))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(refresh.status(), axum::http::StatusCode::OK);
    let body = axum::body::to_bytes(refresh.into_body(), usize::MAX)
        .await
        .unwrap();
    let refreshed: gateway::authn::AuthResponse = serde_json::from_slice(&body).unwrap();
    assert_ne!(refreshed.refresh_token, created.refresh_token);

    let refresh_again_req = serde_json::to_vec(&gateway::authn::RefreshRequest {
        session_id: created.session_id.clone(),
        refresh_token: created.refresh_token.clone(),
    })
    .unwrap();

    let refresh_again = app1
        .oneshot(
            axum::http::Request::builder()
                .method("POST")
                .uri("/v1/auth/refresh")
                .header("content-type", "application/json")
                .body(axum::body::Body::from(refresh_again_req))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(refresh_again.status(), axum::http::StatusCode::UNAUTHORIZED);

    let stored = state
        .storage
        .refresh_sessions
        .get(&format!("v1/sessions/{}", created.session_id))
        .await
        .unwrap()
        .unwrap();
    let value: serde_json::Value = serde_json::from_slice(&stored.value).unwrap();
    assert_eq!(
        value.get("v").and_then(|v| v.as_u64()).unwrap_or(0),
        u64::from(gateway::storage::SCHEMA_VERSION)
    );
}

async fn build_two_replicas() -> (axum::Router, axum::Router, gateway::AppState) {
    let metrics = gateway::observability::init_metrics_for_tests();
    let routing = gateway::routing::RouterState::new(Arc::new(gateway::routing::FixedSource::new(
        gateway::routing::RoutingConfig::empty(),
    )))
    .await
    .unwrap();
    let storage = gateway::storage::GatewayStorage::new_in_memory();
    let authn = gateway::authn::AuthnConfig::for_tests();

    let state = gateway::AppState {
        metrics,
        routing: routing.clone(),
        storage: storage.clone(),
        authn: authn.clone(),
    };

    let app1 = gateway::app(state.clone());
    let app2 = gateway::app(gateway::AppState {
        metrics: gateway::observability::init_metrics_for_tests(),
        routing,
        storage,
        authn,
    });

    (app1, app2, state)
}