Connect to the proper database while testing Geo migrations

ee/db/geo/migrate/20180402170913_add_missing_on_primary_to_job_artifact_registry..rb

@@ -12,6 +12,6 @@ class AddMissingOnPrimaryToJobArtifactRegistry < ActiveRecord::Migration[4.2]
   end
 
   def down
-    remove_column :job_artifact_registry, :missing_on_primary
+    remove_column :job_artifact_registry, :missing_on_primary if column_exists?(:job_artifact_registry, :missing_on_primary)
   end
 end

ee/spec/migrations/geo/fix_state_column_in_file_registry_spec.rb

@@ -4,7 +4,7 @@ require 'spec_helper'
 
 require_migration!
 
-RSpec.describe FixStateColumnInFileRegistry do
+RSpec.describe FixStateColumnInFileRegistry, :geo do
   let(:registry) { table(:file_registry) }
 
   it 'correctly sets registry state value' do

ee/spec/migrations/geo/fix_state_column_in_lfs_object_registry_spec.rb

@@ -4,7 +4,7 @@ require 'spec_helper'
 
 require_migration!
 
-RSpec.describe FixStateColumnInLfsObjectRegistry do
+RSpec.describe FixStateColumnInLfsObjectRegistry, :geo do
   let(:registry) { table(:lfs_object_registry) }
 
   before do
@@ -19,12 +19,12 @@ RSpec.describe FixStateColumnInLfsObjectRegistry do
     pending_registries = registry.where(state: 0)
     synced_registries  = registry.where(state: 2)
 
-    expect(pending_registries.pluck(:id)).to contain_exactly(1, 2)
-    expect(synced_registries.pluck(:id)).to contain_exactly(4)
+    expect(pending_registries.pluck(:lfs_object_id)).to contain_exactly(1, 2)
+    expect(synced_registries.pluck(:lfs_object_id)).to contain_exactly(4)
 
     migrate!
 
-    expect(pending_registries.pluck(:id)).to contain_exactly(1)
-    expect(synced_registries.pluck(:id)).to contain_exactly(2, 4)
+    expect(pending_registries.pluck(:lfs_object_id)).to contain_exactly(1)
+    expect(synced_registries.pluck(:lfs_object_id)).to contain_exactly(2, 4)
   end
 end

ee/spec/support/helpers/ee/migrations_helpers.rb

@@ -45,12 +45,24 @@ module EE
 
     def with_db_config(&block)
       if geo_migration?
-        ::Geo::TrackingBase.connected_to(database: :geo) { yield }
+        with_added_geo_connection { yield }
       else
         yield
       end
     end
 
+    def with_added_geo_connection
+      with_reestablished_active_record_base(reconnect: true) do
+        reconfigure_db_connection(
+          name: :geo,
+          config_model: Geo::TrackingBase,
+          model: ActiveRecord::Base
+        )
+
+        yield
+      end
+    end
+
     def geo_migration?
       self.class.metadata[:geo]
     end