1 files changed, 51 insertions, 0 deletions
diff --git a/spec/tasks/gitlab/db_rake_spec.rb b/spec/tasks/gitlab/db_rake_spec.rb
index edfdb44022b..d1f4a12d8fc 100644
--- a/spec/tasks/gitlab/db_rake_spec.rb
+++ b/spec/tasks/gitlab/db_rake_spec.rb
@@ -297,6 +297,57 @@ RSpec.describe 'gitlab:db namespace rake task' do
     end
   end
 
+  describe '#migrate_with_instrumentation' do
+    subject { run_rake_task('gitlab:db:migration_testing', "[#{filename}]") }
+
+    let(:ctx) { double('ctx', migrations: all_migrations, schema_migration: double, get_all_versions: existing_versions) }
+    let(:instrumentation) { instance_double(Gitlab::Database::Migrations::Instrumentation, observations: observations) }
+    let(:existing_versions) { [1] }
+    let(:all_migrations) { [double('migration1', version: 1), pending_migration] }
+    let(:pending_migration) { double('migration2', version: 2) }
+    let(:filename) { 'results-file.json'}
+    let(:buffer) { StringIO.new }
+    let(:observations) { %w[some data] }
+
+    before do
+      allow(ActiveRecord::Base.connection).to receive(:migration_context).and_return(ctx)
+      allow(Gitlab::Database::Migrations::Instrumentation).to receive(:new).and_return(instrumentation)
+      allow(ActiveRecord::Migrator).to receive_message_chain('new.run').with(any_args).with(no_args)
+
+      allow(instrumentation).to receive(:observe).and_yield
+
+      allow(File).to receive(:open).with(filename, 'wb+').and_yield(buffer)
+    end
+
+    it 'fails when given no filename argument' do
+      expect { run_rake_task('gitlab:db:migration_testing') }.to raise_error(/specify result_file/)
+    end
+
+    it 'fails when the given file already exists' do
+      expect(File).to receive(:exist?).with(filename).and_return(true)
+
+      expect { subject }.to raise_error(/File exists/)
+    end
+
+    it 'instruments the pending migration' do
+      expect(instrumentation).to receive(:observe).with(2).and_yield
+
+      subject
+    end
+
+    it 'executes the pending migration' do
+      expect(ActiveRecord::Migrator).to receive_message_chain('new.run').with(:up, ctx.migrations, ctx.schema_migration, pending_migration.version).with(no_args)
+
+      subject
+    end
+
+    it 'writes observations out to JSON file' do
+      subject
+
+      expect(buffer.string).to eq(observations.to_json)
+    end
+  end
+
   def run_rake_task(task_name, arguments = '')
     Rake::Task[task_name].reenable
     Rake.application.invoke_task("#{task_name}#{arguments}")