2 files changed, 22 insertions, 6 deletions
diff --git a/spec/tasks/gitlab/db/lock_writes_rake_spec.rb b/spec/tasks/gitlab/db/lock_writes_rake_spec.rb
index d03e15224cb..ebea644bbf0 100644
--- a/spec/tasks/gitlab/db/lock_writes_rake_spec.rb
+++ b/spec/tasks/gitlab/db/lock_writes_rake_spec.rb
@@ -89,6 +89,26 @@ RSpec.describe 'gitlab:db:lock_writes', :silence_stdout, :reestablished_active_r
       end
     end
 
+    context 'when running in dry_run mode' do
+      before do
+        stub_env('DRY_RUN', 'true')
+      end
+
+      it 'allows writes on the main tables on the ci database' do
+        run_rake_task('gitlab:db:lock_writes')
+        expect do
+          ci_connection.execute("delete from projects")
+        end.not_to raise_error
+      end
+
+      it 'allows writes on the ci tables on the main database' do
+        run_rake_task('gitlab:db:lock_writes')
+        expect do
+          main_connection.execute("delete from ci_builds")
+        end.not_to raise_error
+      end
+    end
+
     context 'multiple shared databases' do
       before do
         allow(::Gitlab::Database).to receive(:db_config_share_with).and_return(nil)
diff --git a/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb b/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
index f9ebb985255..e95c2e241a8 100644
--- a/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
+++ b/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
@@ -43,10 +43,6 @@ RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablish
   end
 
   shared_examples 'truncating legacy tables' do
-    before do
-      allow(ENV).to receive(:[]).and_return(nil)
-    end
-
     context 'when tables are not locked for writes' do
       it 'raises an error when trying to truncate the tables' do
         error_message = /is not locked for writes. Run the rake task gitlab:db:lock_writes first/
@@ -97,7 +93,7 @@ RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablish
 
       context 'when running in dry_run mode' do
         before do
-          allow(ENV).to receive(:[]).with("DRY_RUN").and_return("true")
+          stub_env('DRY_RUN', 'true')
         end
 
         it 'does not truncate any tables' do
@@ -115,7 +111,7 @@ RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablish
 
       context 'when passing until_table parameter via environment variable' do
         before do
-          allow(ENV).to receive(:[]).with("UNTIL_TABLE").and_return(legacy_table)
+          stub_env('UNTIL_TABLE', legacy_table)
         end
 
         it 'sends the table name to TablesTruncate' do