Refactor large Rust integration tests

apps/admin-api/tests/integration/auth_rate_limit.rs

@@ -0,0 +1,193 @@
+#![allow(dead_code, unused_imports)]
+
+use super::common::*;
+
+use std::{
+    collections::BTreeMap,
+    env, fmt,
+    sync::Arc,
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+use async_trait::async_trait;
+use axum::{Json, Router, routing::post};
+use crank_core::{
+    ExecutionConfig, HttpMethod, MembershipRole, OperationSecurityLevel, Protocol,
+    ResponseCachePolicy, RestTarget, SecretKind, Target, ToolDescription, WorkspaceId,
+};
+use crank_core::{IdentityError, IdentityProvider, IdentityProviderKind, LoginOutcome};
+use crank_mapping::{MappingRule, MappingSet};
+use crank_registry::PostgresRegistry;
+use crank_runtime::SecretCrypto;
+use crank_schema::{Schema, SchemaKind};
+use serde_json::{Value, json};
+use serial_test::serial;
+use tokio::net::TcpListener;
+
+use admin_api::{
+    app::build_app,
+    auth::{AuthSettings, BootstrapAdminConfig, hash_password},
+    service::{AdminService, AdminServiceBuilder, OperationPayload},
+    state::AppState,
+};
+
+const DEFAULT_WORKSPACE_ID: &str = "ws_default";
+const TEST_AUTH_EMAIL: &str = "owner@crank.local";
+const TEST_AUTH_PASSWORD: &str = "test-password";
+const TEST_PASSWORD_PEPPER: &str = "test-password-pepper";
+const TEST_SESSION_SECRET: &str = "test-session-secret";
+const TEST_MASTER_KEY: &str = "test-master-key";
+
+struct TestServer {
+    base_url: String,
+    shutdown: Option<tokio::sync::oneshot::Sender<()>>,
+    handle: Option<tokio::task::JoinHandle<()>>,
+}
+
+impl Drop for TestServer {
+    fn drop(&mut self) {
+        let shutdown = self.shutdown.take();
+        let handle = self.handle.take();
+
+        tokio::task::block_in_place(|| {
+            if let Some(shutdown) = shutdown {
+                let _ = shutdown.send(());
+            }
+            if let Some(handle) = handle {
+                let _ = tokio::runtime::Handle::current().block_on(handle);
+            }
+        });
+    }
+}
+
+impl fmt::Display for TestServer {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str(&self.base_url)
+    }
+}
+
+impl AsRef<str> for TestServer {
+    fn as_ref(&self) -> &str {
+        &self.base_url
+    }
+}
+
+struct RejectingIdentityProvider;
+
+#[async_trait]
+impl IdentityProvider for RejectingIdentityProvider {
+    fn id(&self) -> &str {
+        "rejecting-test-provider"
+    }
+
+    fn kind(&self) -> IdentityProviderKind {
+        IdentityProviderKind::Password
+    }
+
+    async fn login_password(
+        &self,
+        _payload: crank_core::LoginPayload,
+    ) -> Result<LoginOutcome, IdentityError> {
+        Err(IdentityError::BadCredentials)
+    }
+}
+
+#[tokio::test(flavor = "multi_thread")]
+#[serial]
+async fn rejects_rapid_login_requests_with_429() {
+    let registry = test_registry().await;
+    let storage_root = test_storage_root("login_rate_limit");
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let app = build_app(AppState {
+        service: test_service(
+            registry,
+            storage_root,
+            test_auth_settings(),
+            test_secret_crypto(),
+        ),
+        api_rate_limiter: crank_runtime::RequestRateLimiter::new(
+            crank_runtime::RequestRateLimitConfig::new(1, 1).unwrap(),
+        ),
+    });
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
+    let handle = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                let _ = shutdown_rx.await;
+            })
+            .await
+            .unwrap();
+    });
+
+    let client = reqwest::Client::new();
+    let root_url = format!("http://{address}");
+    let first_response = client
+        .post(format!("{root_url}/api/auth/login"))
+        .header("x-request-id", "req_admin_rate_01")
+        .json(&json!({
+            "email": TEST_AUTH_EMAIL,
+            "password": TEST_AUTH_PASSWORD,
+        }))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(first_response.status(), reqwest::StatusCode::OK);
+
+    let second_response = client
+        .post(format!("{root_url}/api/auth/login"))
+        .header("x-request-id", "req_admin_rate_02")
+        .json(&json!({
+            "email": TEST_AUTH_EMAIL,
+            "password": TEST_AUTH_PASSWORD,
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(
+        second_response.status(),
+        reqwest::StatusCode::TOO_MANY_REQUESTS
+    );
+    assert_eq!(
+        second_response.headers()["x-request-id"].to_str().unwrap(),
+        "req_admin_rate_02"
+    );
+    let payload = second_response.json::<Value>().await.unwrap();
+    assert_eq!(payload["error"]["code"], "rate_limited");
+    let retry_after_ms = payload["error"]["context"]["retry_after_ms"]
+        .as_u64()
+        .unwrap();
+    assert!((1..=1000).contains(&retry_after_ms));
+
+    let _ = shutdown_tx.send(());
+    let _ = handle.await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+#[serial]
+async fn login_uses_identity_provider_when_configured() {
+    let registry = test_registry().await;
+    let storage_root = test_storage_root("identity_provider_login");
+    let service = AdminServiceBuilder::new(
+        registry,
+        storage_root,
+        test_auth_settings(),
+        test_secret_crypto(),
+        crank_runtime::community_default().build(),
+    )
+    .with_identity_provider(Arc::new(RejectingIdentityProvider))
+    .build();
+
+    let error = match service
+        .login(admin_api::service::LoginPayload {
+            email: TEST_AUTH_EMAIL.to_owned(),
+            password: TEST_AUTH_PASSWORD.to_owned(),
+        })
+        .await
+    {
+        Ok(_) => panic!("login should delegate to identity provider"),
+        Err(error) => error,
+    };
+    assert_eq!(error.to_string(), "invalid email or password");
+}