From 27b289c3391b1825c07312154ed0067caa247788 Mon Sep 17 00:00:00 2001
From: Dull Bananas <dull.bananas0@gmail.com>
Date: Mon, 13 May 2024 02:51:05 +0000
Subject: [PATCH] add Options with disable_migrations field for test

---
 crates/db_schema/src/schema_setup.rs | 38 ++++++++++++++++++++++++----
 crates/db_schema/src/utils.rs        |  2 +-
 src/lib.rs                           |  6 +++--
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/crates/db_schema/src/schema_setup.rs b/crates/db_schema/src/schema_setup.rs
index 6e4f05293..2a34d725c 100644
--- a/crates/db_schema/src/schema_setup.rs
+++ b/crates/db_schema/src/schema_setup.rs
@@ -38,8 +38,14 @@ fn get_pending_migrations(conn: &mut PgConnection) -> LemmyResult<Vec<Box<dyn Mi
   )
 }
 
-pub fn run(db_url: &str) -> LemmyResult<()> {
-  // Migrations don't support async connection
+#[derive(Default)]
+pub struct Options {
+  /// Only for testing
+  disable_migrations: bool,
+}
+
+pub fn run(db_url: &str, options: &Options) -> LemmyResult<()> {
+  // Migrations don't support async connection, and this function doesn't need to be async
   let mut conn = PgConnection::establish(db_url).with_context(|| "Error connecting to database")?;
 
   let test_enabled = std::env::var("LEMMY_TEST_MIGRATIONS")
@@ -80,9 +86,14 @@ pub fn run(db_url: &str) -> LemmyResult<()> {
     let pending_migrations = get_pending_migrations(conn)?;
 
     // Drop `r` schema and disable the trigger that prevents the Diesel CLI from running migrations
-    conn.batch_execute(
-      "DROP SCHEMA IF EXISTS r CASCADE; SET LOCAL lemmy.enable_migrations TO 'on';",
-    )?;
+    let enable_migrations = if options.disable_migrations {
+      ""
+    } else {
+      "SET LOCAL lemmy.enable_migrations TO 'on';"
+    };
+    conn.batch_execute(&format!(
+      "DROP SCHEMA IF EXISTS r CASCADE;{enable_migrations}"
+    ))?;
 
     for migration in &pending_migrations {
       let name = migration.name();
@@ -112,3 +123,20 @@ pub fn run(db_url: &str) -> LemmyResult<()> {
 
   Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+  use lemmy_utils::{error::LemmyResult, settings::SETTINGS};
+
+  #[test]
+  fn test_schema_setup() -> LemmyResult<()> {
+    let mut options = super::Options::default();
+    let db_url = SETTINGS.get_database_url();
+
+    // Test the forbid_diesel_cli trigger
+    options.disable_migrations = true;
+    super::run(&db_url, &options).expect_err("forbid_diesel_cli trigger should throw error");
+
+    Ok(())
+  }
+}
diff --git a/crates/db_schema/src/utils.rs b/crates/db_schema/src/utils.rs
index e75d7ac51..b489d0b2d 100644
--- a/crates/db_schema/src/utils.rs
+++ b/crates/db_schema/src/utils.rs
@@ -427,7 +427,7 @@ pub async fn build_db_pool() -> LemmyResult<ActualDbPool> {
     }))
     .build()?;
 
-  crate::schema_setup::run(&db_url)?;
+  crate::schema_setup::run(&db_url, &Default::default())?;
 
   Ok(pool)
 }
diff --git a/src/lib.rs b/src/lib.rs
index f3a192325..5fe9cf392 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -40,7 +40,7 @@ use lemmy_apub::{
   VerifyUrlData,
   FEDERATION_HTTP_FETCH_LIMIT,
 };
-use lemmy_db_schema::{source::secret::Secret, utils::build_db_pool};
+use lemmy_db_schema::{schema_setup, source::secret::Secret, utils::build_db_pool};
 use lemmy_federate::{start_stop_federation_workers_cancellable, Opts};
 use lemmy_routes::{feeds, images, nodeinfo, webfinger};
 use lemmy_utils::{
@@ -129,11 +129,13 @@ pub async fn start_lemmy_server(args: CmdArgs) -> LemmyResult<()> {
   println!("Lemmy v{VERSION}");
 
   if let Some(CmdSubcommand::Migration { subcommand }) = args.subcommand {
+    let mut options = schema_setup::Options::default();
+
     match subcommand {
       MigrationSubcommand::Run => {}
     }
 
-    lemmy_db_schema::schema_setup::run(&SETTINGS.get_database_url())?;
+    schema_setup::run(&SETTINGS.get_database_url(), &options)?;
 
     return Ok(());
   }