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}; use backend::routes; use backend::state::AppState; use backend::storage::ObjectStorage; use diesel::connection::SimpleConnection; use diesel::prelude::*; 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::Serialize; 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, }; 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()); Ok(Self { state, router, storage, }) } pub async fn cleanup(&self) -> Result<()> { let pool = self.state.pool.clone(); 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")? } #[allow(dead_code)] pub fn storage(&self) -> Arc { self.storage.clone() } 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(); self.with_conn(move |conn| { let password_hash = hash_password(&password)?; let user = NewUser { id: Uuid::new_v4(), username, password_hash, role, }; diesel::insert_into(backend::schema::users::table) .values(&user) .execute(conn) .context("failed to insert user")?; Ok(user.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(serde::Deserialize)] struct LoginResponse { access_token: String, } let parsed: LoginResponse = serde_json::from_slice(&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 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"); } 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")) } 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 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_tags, document_versions, documents, folders, tags, users RESTART IDENTITY CASCADE;", ) .context("failed to truncate tables")?; Ok(()) } 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()) }