Extend require_migration helper on EE

ee/spec/support/helpers/ee/require_migration.rb

+# frozen_string_literal: true
+
+module EE
+  module RequireMigration
+    extend ActiveSupport::Concern
+
+    MIGRATION_FOLDERS = %w[ee/db/geo/migrate ee/db/geo/post_migrate].freeze
+
+    class_methods do
+      def migration_folders
+        @migration_folders ||= super + MIGRATION_FOLDERS
+      end
+    end
+  end
+end

spec/support/helpers/require_migration.rb

@@ -14,8 +14,7 @@ class RequireMigration
     end
   end
 
-  FOSS_MIGRATION_FOLDERS = %w[db/migrate db/post_migrate].freeze
-  ALL_MIGRATION_FOLDERS = (FOSS_MIGRATION_FOLDERS + %w[ee/db/geo/migrate ee/db/geo/post_migrate]).freeze
+  MIGRATION_FOLDERS = %w[db/migrate db/post_migrate].freeze
   SPEC_FILE_PATTERN = /.+\/(?<file_name>.+)_spec\.rb/.freeze
 
   class << self
@@ -37,11 +36,13 @@ class RequireMigration
     private
 
     def migration_folders
-      Gitlab.ee? ? ALL_MIGRATION_FOLDERS : FOSS_MIGRATION_FOLDERS
+      MIGRATION_FOLDERS
     end
   end
 end
 
+RequireMigration.prepend_if_ee('EE::RequireMigration')
+
 def require_migration!(file_name = nil)
   location_info = caller_locations.first.path.match(RequireMigration::SPEC_FILE_PATTERN)
   file_name ||= location_info[:file_name]