1 files changed, 30 insertions, 1 deletions
diff --git a/spec/tasks/gitlab/storage_rake_spec.rb b/spec/tasks/gitlab/storage_rake_spec.rb
index abd44adbdcc..fefcd05af3b 100644
--- a/spec/tasks/gitlab/storage_rake_spec.rb
+++ b/spec/tasks/gitlab/storage_rake_spec.rb
@@ -2,7 +2,7 @@
 
 require 'rake_helper'
 
-RSpec.describe 'rake gitlab:storage:*' do
+RSpec.describe 'rake gitlab:storage:*', :silence_stdout do
   before do
     Rake.application.rake_require 'tasks/gitlab/storage'
 
@@ -88,6 +88,27 @@ RSpec.describe 'rake gitlab:storage:*' do
     end
   end
 
+  shared_examples 'wait until database is ready' do
+    it 'checks if the database is ready once' do
+      expect(Gitlab::Database).to receive(:exists?).once
+
+      run_rake_task(task)
+    end
+
+    context 'handles custom env vars' do
+      before do
+        stub_env('MAX_DATABASE_CONNECTION_CHECKS' => 3)
+        stub_env('MAX_DATABASE_CONNECTION_INTERVAL' => 0.1)
+      end
+
+      it 'tries for 3 times, polling every 0.1 seconds' do
+        expect(Gitlab::Database).to receive(:exists?).exactly(3).times.and_return(false)
+
+        run_rake_task(task)
+      end
+    end
+  end
+
   describe 'gitlab:storage:migrate_to_hashed' do
     let(:task) { 'gitlab:storage:migrate_to_hashed' }
 
@@ -198,6 +219,10 @@ RSpec.describe 'rake gitlab:storage:*' do
       let(:task) { 'gitlab:storage:legacy_projects' }
       let(:create_collection) { create_list(:project, 3, :legacy_storage) }
     end
+
+    it_behaves_like 'wait until database is ready' do
+      let(:task) { 'gitlab:storage:legacy_projects' }
+    end
   end
 
   describe 'gitlab:storage:list_legacy_projects' do
@@ -227,6 +252,10 @@ RSpec.describe 'rake gitlab:storage:*' do
       let(:project) { create(:project, storage_version: 1) }
       let(:create_collection) { create_list(:upload, 3, model: project) }
     end
+
+    it_behaves_like 'wait until database is ready' do
+      let(:task) { 'gitlab:storage:legacy_attachments' }
+    end
   end
 
   describe 'gitlab:storage:list_legacy_attachments' do