From b5d610a62970fb2e8dab8f919e8757eb92dd2593 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Sat, 10 Feb 2024 19:58:30 +0200
Subject: [PATCH] foo

---
 .../migration_helpers.rb                      | 21 +++++++++++++++++++
 test/schema_statements/misc_test.rb           |  5 +++++
 2 files changed, 26 insertions(+)

diff --git a/lib/online_migrations/background_migrations/migration_helpers.rb b/lib/online_migrations/background_migrations/migration_helpers.rb
index 7aa3d12..7b7c96d 100644
--- a/lib/online_migrations/background_migrations/migration_helpers.rb
+++ b/lib/online_migrations/background_migrations/migration_helpers.rb
@@ -381,6 +381,27 @@ def enqueue_background_migration(migration_name, *arguments, **options)
         migration
       end
 
+      def ensure_background_migration_succeeded(migration_name, arguments: nil)
+        migration_name = migration_name.name if migration_name.is_a?(Class)
+
+        configuration = { migration_name: migration_name }
+
+        if arguments
+          arguments = Array(arguments)
+          migration = Migration.for_migration_name(migration_name).first
+          configuration[:arguments] = arguments.to_json
+        else
+          migration = Migration.for_migration_name(migration_name).first
+        end
+
+        if migration.nil?
+          Utils.say("Could not find background migration for the given configuration: #{configuration}")
+        elsif !migration.succeeded?
+          raise "Expected background migration for the given configuration to be marked as 'succeeded', "\
+                "but it is '#{migration.status}': #{configuration}"
+        end
+      end
+
       # @private
       def create_background_migration(migration_name, *arguments, **options)
         options.assert_valid_keys(:batch_column_name, :min_value, :max_value, :batch_size, :sub_batch_size,
diff --git a/test/schema_statements/misc_test.rb b/test/schema_statements/misc_test.rb
index a3e1270..1c54ab4 100644
--- a/test/schema_statements/misc_test.rb
+++ b/test/schema_statements/misc_test.rb
@@ -203,6 +203,11 @@ def test_run_background_migrations_inline_configured_to_custom_proc
       assert_nil user.reload.admin
     end
 
+    def test_ensure_background_migration_succeeded
+      m = @connection.enqueue_background_migration("MakeAllNonAdmins")
+      assert m.succeeded?
+    end
+
     def test_disable_statement_timeout
       prev_value = get_statement_timeout
       set_statement_timeout(10)