Fixing tests and re-locating error checking.

2 files changed, 10 insertions, 9 deletions

diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb
index e06a581094..f5753fe426 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -164,7 +164,7 @@ module ActiveRecord
   end
 
   class EnvironmentMismatchError < ActiveRecordError
-    def initialize(current: , stored: )
+    def initialize(current: nil, stored: nil)
       msg =  "You are attempting to modify a database that was last run in #{ stored } environment.\n"
       msg << "You are running in #{ current } environment."
       msg << "if you are sure you want to continue, run the same command with the environment variable\n"
@@ -1238,7 +1238,12 @@ module ActiveRecord
     end
 
     def self.last_stored_environment
-      ActiveRecord::InternalMetadata[:environment]
+      return nil if current_version == 0
+      raise NoEnvironmentInSchemaError unless ActiveRecord::InternalMetadata.table_exists?
+
+      environment = ActiveRecord::InternalMetadata[:environment]
+      raise NoEnvironmentInSchemaError unless environment
+      environment
     end
 
     def self.current_environment
@@ -1246,11 +1251,7 @@ module ActiveRecord
     end
 
     def self.protected_environment?
-      return false if current_version == 0
-      raise NoEnvironmentInSchemaError unless ActiveRecord::InternalMetadata.table_exists?
-
-      raise NoEnvironmentInSchemaError unless last_stored_environment
-      ActiveRecord::Base.protected_environments.include?(last_stored_environment)
+      ActiveRecord::Base.protected_environments.include?(last_stored_environment) if last_stored_environment
     end
 
     def up?
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index d51f38c93f..7b8f27699a 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -51,8 +51,8 @@ module ActiveRecord
             raise ActiveRecord::ProtectedEnvironmentError.new(stored)
           end
 
-          if current != stored
-            raise EnvironmentMismatchError.new(current: current, stored: stored)
+          if stored && stored != current
+            raise ActiveRecord::EnvironmentMismatchError.new(current: current, stored: stored)
           end
         end
       rescue ActiveRecord::NoDatabaseError