foo

backend/tests/common/mod.rs

@@ -11,12 +11,13 @@ use axum::Router;
 use backend::auth::jwt::JwtService;
 use backend::config::AppConfig;
 use backend::db::{self, PgPool};
-use backend::models::{Job, NewUser, NewUserMembership};
+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;
@@ -151,11 +152,15 @@ impl TestApp {
         let state = AppState::new(pool.clone(), config, storage_for_state, jwt);
         let router = routes::create_router(state.clone());
 
-        Ok(Self {
+        let app = Self {
             state,
             router,
             storage,
-        })
+        };
+
+        app.ensure_default_tenant().await?;
+
+        Ok(app)
     }
 
     pub async fn cleanup(&self) -> Result<()> {
@@ -168,7 +173,10 @@ impl TestApp {
             Ok(())
         })
         .await
-        .context("cleanup task panicked")?
+        .context("cleanup task panicked")?;
+
+        self.ensure_default_tenant().await?;
+        Ok(())
     }
 
     #[allow(dead_code)]
@@ -176,6 +184,19 @@ impl TestApp {
         self.storage.clone()
     }
 
+    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();
@@ -184,7 +205,7 @@ impl TestApp {
             .state
             .tenants
             .tenant_id_for_slug(&self.state.config.default_tenant_slug)
-            .context("default tenant not found")?;
+            .map_err(|err| anyhow!("default tenant not found: {:?}", err))?;
         self.with_conn(move |conn| {
             let password_hash = hash_password(&password)?;
             let user = NewUser {
@@ -213,6 +234,60 @@ impl TestApp {
         .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),
+                    ))
+                    .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> {
@@ -490,7 +565,22 @@ async fn prepare_database(pool: &PgPool) -> Result<()> {
 
 fn truncate_all(conn: &mut PgConnection) -> Result<()> {
     conn.batch_execute(
-        "TRUNCATE TABLE document_tags, document_versions, documents, folders, tags, users RESTART IDENTITY CASCADE;",
+        "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(())