use std::collections::HashMap; use std::env; use std::sync::Arc; use std::time::Duration; use anyhow::{anyhow, ensure, Context, Result}; use async_trait::async_trait; use axum::body::Body; use axum::http::{Method, Request, StatusCode}; use axum::Router; use backend::auth::jwt::JwtService; use backend::config::AppConfig; use backend::db::{self, PgPool}; use backend::models::{Job, NewUser, NewUserMembership, Tenant}; use backend::routes; use backend::state::AppState; use backend::storage::ObjectStorage; use diesel::connection::SimpleConnection; use diesel::prelude::*; use diesel::OptionalExtension; use diesel::PgConnection; use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use http_body_util::BodyExt; use once_cell::sync::Lazy; use rand::rngs::OsRng; use serde::{Deserialize, Serialize}; use serde_json; use tokio::sync::Mutex; use tower::util::ServiceExt; use uuid::Uuid; const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations"); static DB_LOCK: Lazy> = Lazy::new(|| Mutex::new(())); #[allow(dead_code)] #[derive(Clone)] pub struct StoredObject { pub key: String, pub bytes: Vec, pub content_type: Option, pub content_disposition: Option, } #[derive(Default)] pub struct FakeStorage { objects: Mutex>, } #[async_trait] impl ObjectStorage for FakeStorage { async fn put_object( &self, key: &str, bytes: Vec, content_type: Option, content_disposition: Option, ) -> Result<()> { let stored = StoredObject { key: key.to_string(), bytes, content_type, content_disposition, }; let mut guard = self.objects.lock().await; guard.insert(stored.key.clone(), stored); Ok(()) } async fn presign_get_object(&self, key: &str, expires_in: Duration) -> Result { let guard = self.objects.lock().await; ensure!(guard.contains_key(key), "object {key} missing"); Ok(format!( "https://fake-storage/{key}?expires_in={}", expires_in.as_secs() )) } async fn get_object(&self, key: &str) -> Result> { let guard = self.objects.lock().await; guard .get(key) .map(|obj| obj.bytes.clone()) .ok_or_else(|| anyhow!("object {key} missing")) } async fn delete_object(&self, key: &str) -> Result<()> { let mut guard = self.objects.lock().await; guard.remove(key); Ok(()) } } impl FakeStorage { #[allow(dead_code)] pub async fn get(&self, key: &str) -> Option { let guard = self.objects.lock().await; guard.get(key).cloned() } #[allow(dead_code)] pub async fn object_count(&self) -> usize { let guard = self.objects.lock().await; guard.len() } } pub struct TestApp { pub state: AppState, router: Router, storage: Arc, } impl TestApp { pub async fn new() -> Result { let database_url = env::var("TEST_DATABASE_URL") .context("TEST_DATABASE_URL must be set for integration tests")?; let config = AppConfig { database_url: database_url.clone(), database_max_pool_size: db::DEFAULT_MAX_POOL_SIZE, server_host: "127.0.0.1".to_string(), server_port: 0, webdav_host: "127.0.0.1".to_string(), webdav_port: 0, jwt_secret: "test-secret".to_string(), jwt_issuer: "test-issuer".to_string(), jwt_audience: "test-audience".to_string(), jwt_expiry_minutes: 60, download_token_audience: "test-download".to_string(), download_token_expiry_minutes: 60, refresh_token_expiry_days: 30, refresh_cookie_secure: false, refresh_cookie_domain: None, cors_allowed_origin: None, aws_endpoint_url: None, aws_access_key_id: None, aws_secret_access_key: None, aws_region: "us-east-1".to_string(), s3_bucket: "test-bucket".to_string(), quickwit_endpoint: None, quickwit_index: None, default_tenant_slug: "admin".to_string(), }; let pool = db::init_pool_with_size(&config.database_url, config.database_max_pool_size)?; prepare_database(&pool).await?; let storage = Arc::new(FakeStorage::default()); let storage_for_state: Arc = storage.clone(); let jwt = JwtService::from_config(&config)?; let state = AppState::new(pool.clone(), config, storage_for_state, jwt); let router = routes::create_router(state.clone()); let app = Self { state, router, storage, }; app.ensure_default_tenant().await?; Ok(app) } pub async fn cleanup(&self) -> Result<()> { let pool = self.state.pool.clone(); let _ = tokio::task::spawn_blocking(move || -> Result<()> { let mut conn = pool .get() .map_err(|err| anyhow!("failed to get cleanup connection: {err}"))?; truncate_all(&mut conn)?; Ok(()) }) .await .context("cleanup task panicked")?; self.ensure_default_tenant().await?; Ok(()) } #[allow(dead_code)] pub fn storage(&self) -> Arc { self.storage.clone() } #[allow(dead_code)] pub async fn storage_key_for(&self, key: &str) -> Result { let tenant = self .state .tenants .get_by_slug(&self.state.config.default_tenant_slug) .map_err(|err| anyhow!("default tenant not found: {:?}", err))?; let root = tenant .storage_root .clone() .ok_or_else(|| anyhow!("default tenant missing storage root"))?; Ok(format!("{}{}", root, key)) } pub async fn insert_user(&self, username: &str, password: &str, role: &str) -> Result { let username = username.to_string(); let password = password.to_string(); let role = role.to_string(); let tenant_id = self .state .tenants .tenant_id_for_slug(&self.state.config.default_tenant_slug) .map_err(|err| anyhow!("default tenant not found: {:?}", err))?; self.with_conn(move |conn| { let password_hash = hash_password(&password)?; let user = NewUser { id: Uuid::new_v4(), username, password_hash, }; diesel::insert_into(backend::schema::users::table) .values(&user) .execute(conn) .context("failed to insert user")?; let membership = NewUserMembership { id: Uuid::new_v4(), user_id: user.id, tenant_id, role, }; diesel::insert_into(backend::schema::user_memberships::table) .values(&membership) .execute(conn) .context("failed to insert user membership")?; Ok(user.id) }) .await } async fn ensure_default_tenant(&self) -> Result { let slug_value = self.state.config.default_tenant_slug.clone(); let quickwit_enabled = self.state.config.quickwit_endpoint.is_some(); self.with_conn(move |conn| { use backend::schema::tenants::dsl as tenants_dsl; let existing = tenants_dsl::tenants .filter(tenants_dsl::slug.eq(&slug_value)) .first::(conn) .optional() .context("failed to load default tenant")?; let tenant_id = if let Some(current) = existing { let desired_root = current .storage_root .clone() .filter(|root| root.ends_with('/')) .unwrap_or_else(|| format!("test-tenants/{}/", current.id)); if current.storage_root.as_deref() != Some(desired_root.as_str()) { diesel::update(tenants_dsl::tenants.filter(tenants_dsl::id.eq(current.id))) .set(tenants_dsl::storage_root.eq(Some(desired_root))) .execute(conn) .context("failed to update default tenant storage root")?; } current.id } else { let new_id = Uuid::new_v4(); let root = format!("test-tenants/{}/", new_id); let quickwit_value = if quickwit_enabled { Some(format!("documents-{}", new_id)) } else { None }; diesel::insert_into(tenants_dsl::tenants) .values(( tenants_dsl::id.eq(new_id), tenants_dsl::slug.eq(&slug_value), tenants_dsl::storage_root.eq(Some(root)), tenants_dsl::quickwit_index.eq(quickwit_value), )) .execute(conn) .context("failed to insert default tenant")?; new_id }; Ok(tenant_id) }) .await } pub async fn login_token(&self, username: &str, password: &str) -> Result { #[derive(Serialize)] struct LoginPayload<'a> { username: &'a str, password: &'a str, } let response = self .post_json( "/api/auth/login", &LoginPayload { username, password }, None, ) .await?; ensure!( response.status() == StatusCode::OK, "login failed with status {}", response.status() ); let body = body_to_vec(response.into_body()).await?; #[derive(Deserialize)] struct LoginResponse { access_token: String, } if let Ok(parsed) = serde_json::from_slice::(&body) { return Ok(parsed.access_token); } #[derive(Deserialize)] struct TenantSummary { tenant_id: Uuid, _slug: String, } #[derive(Deserialize)] struct TenantSelectionResponse { selection_token: String, tenants: Vec, } let selection: TenantSelectionResponse = serde_json::from_slice(&body)?; ensure!( !selection.tenants.is_empty(), "login returned no tenant options", ); #[derive(Serialize)] struct SelectTenantPayload { tenant_id: Uuid, } let target_tenant = selection.tenants[0].tenant_id; let select_response = self .post_json( "/api/auth/select-tenant", &SelectTenantPayload { tenant_id: target_tenant, }, Some(&selection.selection_token), ) .await?; ensure!( select_response.status() == StatusCode::OK, "tenant selection failed with status {}", select_response.status() ); let select_body = body_to_vec(select_response.into_body()).await?; let parsed: LoginResponse = serde_json::from_slice(&select_body)?; Ok(parsed.access_token) } #[allow(dead_code)] pub async fn clear_jobs(&self) -> Result<()> { self.with_conn(|conn| { use backend::schema::jobs::dsl::jobs as jobs_table; diesel::delete(jobs_table) .execute(conn) .context("failed to clear jobs")?; Ok(()) }) .await } #[allow(dead_code)] pub async fn jobs_by_type(&self, ty: &str) -> Result> { let ty = ty.to_string(); self.with_conn(move |conn| { use backend::schema::jobs::dsl::{job_type as job_type_col, jobs as jobs_table}; let rows = jobs_table .filter(job_type_col.eq(&ty)) .load::(conn) .context("failed to load jobs")?; Ok(rows) }) .await } pub async fn post_json( &self, path: &str, payload: &T, token: Option<&str>, ) -> Result> { let body = serde_json::to_vec(payload)?; let mut builder = Request::builder() .method(Method::POST) .uri(path) .header("content-type", "application/json"); if let Some(token) = token { builder = builder.header("authorization", format!("Bearer {token}")); } let request = builder.body(Body::from(body))?; Ok(self .router .clone() .oneshot(request) .await .expect("infallible response")) } #[allow(dead_code)] pub async fn patch_json( &self, path: &str, payload: &T, token: Option<&str>, ) -> Result> { let body = serde_json::to_vec(payload)?; let mut builder = Request::builder() .method(Method::PATCH) .uri(path) .header("content-type", "application/json"); if let Some(token) = token { builder = builder.header("authorization", format!("Bearer {token}")); } let request = builder.body(Body::from(body))?; Ok(self .router .clone() .oneshot(request) .await .expect("infallible response")) } pub async fn get(&self, path: &str, token: Option<&str>) -> Result> { let mut builder = Request::builder().method(Method::GET).uri(path); if let Some(token) = token { builder = builder.header("authorization", format!("Bearer {token}")); } let request = builder.body(Body::empty())?; Ok(self .router .clone() .oneshot(request) .await .expect("infallible response")) } #[allow(dead_code)] pub async fn delete(&self, path: &str, token: Option<&str>) -> Result> { let builder = Request::builder().method(Method::DELETE).uri(path); let builder = if let Some(token) = token { builder.header("authorization", format!("Bearer {token}")) } else { builder }; let request = builder.body(Body::empty())?; Ok(self .router .clone() .oneshot(request) .await .expect("infallible response")) } #[allow(dead_code)] pub async fn upload_document( &self, path: &str, filename: &str, content_type: &str, data: &[u8], folder_id: Option, token: &str, ) -> Result> { let extras = UploadExtras::empty(); self.upload_document_with_extras( path, filename, content_type, data, folder_id, extras, token, ) .await } pub async fn upload_document_with_options( &self, path: &str, filename: &str, content_type: &str, data: &[u8], folder_id: Option, title: Option<&str>, metadata_json: Option<&str>, token: &str, ) -> Result> { let extras = UploadExtras { title, metadata_json, tag_ids_json: None, correspondents_json: None, issued_at: None, skip_existing: false, }; self.upload_document_with_extras( path, filename, content_type, data, folder_id, extras, token, ) .await } pub async fn upload_document_with_extras( &self, path: &str, filename: &str, content_type: &str, data: &[u8], folder_id: Option, extras: UploadExtras<'_>, token: &str, ) -> Result> { let boundary = format!("boundary-{}", Uuid::new_v4()); let mut body = Vec::new(); body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend( format!( "Content-Disposition: form-data; name=\"file\"; filename=\"{}\"\r\n", filename ) .as_bytes(), ); body.extend(format!("Content-Type: {}\r\n\r\n", content_type).as_bytes()); body.extend(data); body.extend(b"\r\n"); if let Some(folder) = folder_id { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"folder_id\"\r\n\r\n"); body.extend(folder.to_string().as_bytes()); body.extend(b"\r\n"); } if let Some(title_value) = extras.title { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"title\"\r\n\r\n"); body.extend(title_value.as_bytes()); body.extend(b"\r\n"); } if let Some(metadata_value) = extras.metadata_json { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"metadata\"\r\n\r\n"); body.extend(metadata_value.as_bytes()); body.extend(b"\r\n"); } if let Some(tag_ids_value) = extras.tag_ids_json { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"tag_ids\"\r\n\r\n"); body.extend(tag_ids_value.as_bytes()); body.extend(b"\r\n"); } if let Some(correspondents_value) = extras.correspondents_json { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"correspondents\"\r\n\r\n"); body.extend(correspondents_value.as_bytes()); body.extend(b"\r\n"); } if let Some(issued_at_value) = extras.issued_at { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"issued_at\"\r\n\r\n"); body.extend(issued_at_value.as_bytes()); body.extend(b"\r\n"); } if extras.skip_existing { body.extend(format!("--{boundary}\r\n").as_bytes()); body.extend(b"Content-Disposition: form-data; name=\"skip_existing\"\r\n\r\ntrue\r\n"); } body.extend(format!("--{boundary}--\r\n").as_bytes()); let builder = Request::builder() .method(Method::POST) .uri(path) .header( "content-type", format!("multipart/form-data; boundary={boundary}"), ) .header("authorization", format!("Bearer {token}")); let request = builder.body(Body::from(body))?; Ok(self .router .clone() .oneshot(request) .await .expect("infallible response")) } pub async fn with_conn(&self, f: F) -> Result where F: FnOnce(&mut PgConnection) -> Result + Send + 'static, T: Send + 'static, { let pool = self.state.pool.clone(); tokio::task::spawn_blocking(move || { let mut conn = pool .get() .map_err(|err| anyhow!("failed to get database connection: {err}"))?; f(&mut conn) }) .await .context("connection task panicked")? } } pub struct UploadExtras<'a> { pub title: Option<&'a str>, pub metadata_json: Option<&'a str>, pub tag_ids_json: Option<&'a str>, pub correspondents_json: Option<&'a str>, pub issued_at: Option<&'a str>, pub skip_existing: bool, } impl<'a> UploadExtras<'a> { pub fn empty() -> Self { Self { title: None, metadata_json: None, tag_ids_json: None, correspondents_json: None, issued_at: None, skip_existing: false, } } } pub async fn acquire_db_lock() -> tokio::sync::MutexGuard<'static, ()> { DB_LOCK.lock().await } pub async fn body_to_vec(body: Body) -> Result> { let collected = body .collect() .await .map_err(|err| anyhow!("failed to read response body: {err}"))?; Ok(collected.to_bytes().to_vec()) } async fn prepare_database(pool: &PgPool) -> Result<()> { let pool = pool.clone(); tokio::task::spawn_blocking(move || -> Result<()> { let mut conn = pool .get() .map_err(|err| anyhow!("failed to acquire connection: {err}"))?; conn.run_pending_migrations(MIGRATIONS) .map_err(|err| anyhow!("failed to run migrations: {err}"))?; truncate_all(&mut conn)?; Ok(()) }) .await .context("migration task panicked")? } fn truncate_all(conn: &mut PgConnection) -> Result<()> { conn.batch_execute( "TRUNCATE TABLE \ document_asset_objects, \ document_assets, \ document_correspondents, \ correspondents, \ document_tags, \ document_versions, \ documents, \ folders, \ jobs, \ refresh_tokens, \ tags, \ user_memberships, \ users, \ tenants \ RESTART IDENTITY CASCADE;", ) .context("failed to truncate tables")?; Ok(()) } pub fn hash_password(password: &str) -> Result { use argon2::password_hash::{PasswordHasher, SaltString}; use argon2::Argon2; let salt = SaltString::generate(&mut OsRng); Ok(Argon2::default() .hash_password(password.as_bytes(), &salt) .map_err(|err| anyhow!("failed to hash password: {err}"))? .to_string()) }