1 files changed, 22 insertions, 2 deletions
diff --git a/spec/support/helpers/require_migration.rb b/spec/support/helpers/require_migration.rb
index d3f192a4142..c2902aa4ec7 100644
--- a/spec/support/helpers/require_migration.rb
+++ b/spec/support/helpers/require_migration.rb
@@ -3,26 +3,46 @@
 require 'find'
 
 class RequireMigration
-  MIGRATION_FOLDERS = %w(db/migrate db/post_migrate ee/db/geo/migrate ee/db/geo/post_migrate).freeze
+  class AutoLoadError < RuntimeError
+    MESSAGE = "Can not find any migration file for `%{file_name}`!\n" \
+              "You can try to provide the migration file name manually."
+
+    def initialize(file_name)
+      message = format(MESSAGE, file_name: file_name)
+
+      super(message)
+    end
+  end
+
+  MIGRATION_FOLDERS = %w[db/migrate db/post_migrate].freeze
   SPEC_FILE_PATTERN = /.+\/(?<file_name>.+)_spec\.rb/.freeze
 
   class << self
     def require_migration!(file_name)
       file_paths = search_migration_file(file_name)
+      raise AutoLoadError.new(file_name) unless file_paths.first
 
       require file_paths.first
     end
 
     def search_migration_file(file_name)
-      MIGRATION_FOLDERS.flat_map do |path|
+      migration_folders.flat_map do |path|
         migration_path = Rails.root.join(path).to_s
 
         Find.find(migration_path).grep(/\d+_#{file_name}\.rb/)
       end
     end
+
+    private
+
+    def 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]