papercrate/backend/tests/common/mod.rs

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::{header, 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, TenantStatus};
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<Mutex<()>> = Lazy::new(|| Mutex::new(()));

#[allow(dead_code)]
#[derive(Clone)]
pub struct StoredObject {
    pub key: String,
    pub bytes: Vec<u8>,
    pub content_type: Option<String>,
    pub content_disposition: Option<String>,
}

#[derive(Default)]
pub struct FakeStorage {
    objects: Mutex<HashMap<String, StoredObject>>,
}

#[async_trait]
impl ObjectStorage for FakeStorage {
    async fn put_object(
        &self,
        key: &str,
        bytes: Vec<u8>,
        content_type: Option<String>,
        content_disposition: Option<String>,
    ) -> 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<String> {
        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<Vec<u8>> {
        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<StoredObject> {
        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<FakeStorage>,
}

impl TestApp {
    pub async fn new() -> Result<Self> {
        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(),
            worker_max_document_bytes: 200 * 1024 * 1024,
        };

        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<dyn ObjectStorage> = 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<FakeStorage> {
        self.storage.clone()
    }

    #[allow(dead_code)]
    pub async fn storage_key_for(&self, key: &str) -> Result<String> {
        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<Uuid> {
        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<Uuid> {
        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::<Tenant>(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),
                        tenants_dsl::status.eq(TenantStatus::Active),
                    ))
                    .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<String> {
        #[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::<LoginResponse>(&body) {
            return Ok(parsed.access_token);
        }

        #[derive(Deserialize)]
        struct TenantSummary {
            id: Uuid,
            _slug: String,
        }

        #[derive(Deserialize)]
        struct TenantSelectionResponse {
            access_token: String,
            tenants: Vec<TenantSummary>,
        }

        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].id;
        let select_response = self
            .post_json(
                "/api/auth/select-tenant",
                &SelectTenantPayload {
                    tenant_id: target_tenant,
                },
                Some(&selection.access_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<Vec<Job>> {
        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::<Job>(conn)
                .context("failed to load jobs")?;
            Ok(rows)
        })
        .await
    }

    pub async fn post_json<T: Serialize + ?Sized>(
        &self,
        path: &str,
        payload: &T,
        token: Option<&str>,
    ) -> Result<hyper::Response<Body>> {
        self.post_json_with_cookie(path, payload, token, None).await
    }

    pub async fn post_json_with_cookie<T: Serialize + ?Sized>(
        &self,
        path: &str,
        payload: &T,
        token: Option<&str>,
        cookie: Option<&str>,
    ) -> Result<hyper::Response<Body>> {
        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}"));
        }
        if let Some(cookie) = cookie {
            builder = builder.header(header::COOKIE, cookie);
        }
        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<T: Serialize + ?Sized>(
        &self,
        path: &str,
        payload: &T,
        token: Option<&str>,
    ) -> Result<hyper::Response<Body>> {
        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<hyper::Response<Body>> {
        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<hyper::Response<Body>> {
        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<Uuid>,
        token: &str,
    ) -> Result<hyper::Response<Body>> {
        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<Uuid>,
        title: Option<&str>,
        metadata_json: Option<&str>,
        token: &str,
    ) -> Result<hyper::Response<Body>> {
        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<Uuid>,
        extras: UploadExtras<'_>,
        token: &str,
    ) -> Result<hyper::Response<Body>> {
        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<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut PgConnection) -> Result<T> + 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<Vec<u8>> {
    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, \
            webdav_tokens, \
            user_memberships, \
            users, \
            tenants \
        RESTART IDENTITY CASCADE;",
    )
    .context("failed to truncate tables")?;
    Ok(())
}

pub fn hash_password(password: &str) -> Result<String> {
    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())
}