Add latest changes from gitlab-org/gitlab@master

8 files changed, 134 insertions, 80 deletions

diff --git a/spec/controllers/admin/groups_controller_spec.rb b/spec/controllers/admin/groups_controller_spec.rb
index 8e31ef12adf..d9b7e00fd75 100644
--- a/spec/controllers/admin/groups_controller_spec.rb
+++ b/spec/controllers/admin/groups_controller_spec.rb
@@ -3,9 +3,9 @@
 require 'spec_helper'
 
 RSpec.describe Admin::GroupsController do
-  let(:group) { create(:group) }
-  let(:project) { create(:project, namespace: group) }
-  let(:admin) { create(:admin) }
+  let_it_be(:group) { create(:group) }
+  let_it_be(:project) { create(:project, namespace: group) }
+  let_it_be(:admin) { create(:admin) }
 
   before do
     sign_in(admin)
@@ -46,9 +46,9 @@ RSpec.describe Admin::GroupsController do
   end
 
   describe 'PUT #members_update' do
-    let(:group_user) { create(:user) }
+    let_it_be(:group_user) { create(:user) }
 
-    it 'adds user to members' do
+    it 'adds user to members', :aggregate_failures, :snowplow do
       put :members_update, params: {
                              id: group,
                              user_ids: group_user.id,
@@ -58,9 +58,16 @@ RSpec.describe Admin::GroupsController do
       expect(controller).to set_flash.to 'Users were successfully added.'
       expect(response).to redirect_to(admin_group_path(group))
       expect(group.users).to include group_user
+      expect_snowplow_event(
+        category: 'Members::CreateService',
+        action: 'create_member',
+        label: 'admin-group-page',
+        property: 'existing_user',
+        user: admin
+      )
     end
 
-    it 'can add unlimited members' do
+    it 'can add unlimited members', :aggregate_failures do
       put :members_update, params: {
                              id: group,
                              user_ids: 1.upto(1000).to_a.join(','),
@@ -71,7 +78,7 @@ RSpec.describe Admin::GroupsController do
       expect(response).to redirect_to(admin_group_path(group))
     end
 
-    it 'adds no user to members' do
+    it 'adds no user to members', :aggregate_failures do
       put :members_update, params: {
                              id: group,
                              user_ids: '',
diff --git a/spec/controllers/groups/group_members_controller_spec.rb b/spec/controllers/groups/group_members_controller_spec.rb
index a3e0c9f9ca4..4b17326de09 100644
--- a/spec/controllers/groups/group_members_controller_spec.rb
+++ b/spec/controllers/groups/group_members_controller_spec.rb
@@ -17,7 +17,7 @@ RSpec.describe Groups::GroupMembersController do
   end
 
   describe 'GET index' do
-    it 'renders index with 200 status code' do
+    it 'renders index with 200 status code', :aggregate_failures do
       get :index, params: { group_id: group }
 
       expect(response).to have_gitlab_http_status(:ok)
@@ -118,7 +118,7 @@ RSpec.describe Groups::GroupMembersController do
         group.add_developer(user)
       end
 
-      it 'returns 403' do
+      it 'returns 403', :aggregate_failures do
         post :create, params: {
                         group_id: group,
                         user_ids: group_user.id,
@@ -135,7 +135,7 @@ RSpec.describe Groups::GroupMembersController do
         group.add_owner(user)
       end
 
-      it 'adds user to members' do
+      it 'adds user to members', :aggregate_failures, :snowplow do
         post :create, params: {
                         group_id: group,
                         user_ids: group_user.id,
@@ -145,9 +145,16 @@ RSpec.describe Groups::GroupMembersController do
         expect(controller).to set_flash.to 'Users were successfully added.'
         expect(response).to redirect_to(group_group_members_path(group))
         expect(group.users).to include group_user
+        expect_snowplow_event(
+          category: 'Members::CreateService',
+          action: 'create_member',
+          label: 'group-members-page',
+          property: 'existing_user',
+          user: user
+        )
       end
 
-      it 'adds no user to members' do
+      it 'adds no user to members', :aggregate_failures do
         post :create, params: {
                         group_id: group,
                         user_ids: '',
@@ -177,7 +184,7 @@ RSpec.describe Groups::GroupMembersController do
       context 'when set to a date in the past' do
         let(:expires_at) { 2.days.ago }
 
-        it 'does not add user to members' do
+        it 'does not add user to members', :aggregate_failures do
           subject
 
           expect(flash[:alert]).to include('Expires at cannot be a date in the past')
@@ -189,7 +196,7 @@ RSpec.describe Groups::GroupMembersController do
       context 'when set to a date in the future' do
         let(:expires_at) { 5.days.from_now }
 
-        it 'adds user to members' do
+        it 'adds user to members', :aggregate_failures do
           subject
 
           expect(controller).to set_flash.to 'Users were successfully added.'
@@ -326,7 +333,7 @@ RSpec.describe Groups::GroupMembersController do
           group.add_developer(user)
         end
 
-        it 'returns 403' do
+        it 'returns 403', :aggregate_failures do
           delete :destroy, params: { group_id: group, id: member }
 
           expect(response).to have_gitlab_http_status(:forbidden)
@@ -339,7 +346,7 @@ RSpec.describe Groups::GroupMembersController do
           group.add_owner(user)
         end
 
-        it '[HTML] removes user from members' do
+        it '[HTML] removes user from members', :aggregate_failures do
           delete :destroy, params: { group_id: group, id: member }
 
           expect(controller).to set_flash.to 'User was successfully removed from group.'
@@ -348,7 +355,7 @@ RSpec.describe Groups::GroupMembersController do
           expect(sub_group.members).to include sub_member
         end
 
-        it '[HTML] removes user from members including subgroups and projects' do
+        it '[HTML] removes user from members including subgroups and projects', :aggregate_failures do
           delete :destroy, params: { group_id: group, id: member, remove_sub_memberships: true }
 
           expect(controller).to set_flash.to 'User was successfully removed from group and any subgroups and projects.'
@@ -357,7 +364,7 @@ RSpec.describe Groups::GroupMembersController do
           expect(sub_group.members).not_to include sub_member
         end
 
-        it '[JS] removes user from members' do
+        it '[JS] removes user from members', :aggregate_failures do
           delete :destroy, params: { group_id: group, id: member }, xhr: true
 
           expect(response).to be_successful
@@ -386,7 +393,7 @@ RSpec.describe Groups::GroupMembersController do
           group.add_developer(user)
         end
 
-        it 'removes user from members' do
+        it 'removes user from members', :aggregate_failures do
           delete :leave, params: { group_id: group }
 
           expect(controller).to set_flash.to "You left the \"#{group.name}\" group."
@@ -394,7 +401,7 @@ RSpec.describe Groups::GroupMembersController do
           expect(group.users).not_to include user
         end
 
-        it 'supports json request' do
+        it 'supports json request', :aggregate_failures do
           delete :leave, params: { group_id: group }, format: :json
 
           expect(response).to have_gitlab_http_status(:ok)
@@ -421,7 +428,7 @@ RSpec.describe Groups::GroupMembersController do
           group.request_access(user)
         end
 
-        it 'removes user from members' do
+        it 'removes user from members', :aggregate_failures do
           delete :leave, params: { group_id: group }
 
           expect(controller).to set_flash.to 'Your access request to the group has been withdrawn.'
@@ -438,7 +445,7 @@ RSpec.describe Groups::GroupMembersController do
       sign_in(user)
     end
 
-    it 'creates a new GroupMember that is not a team member' do
+    it 'creates a new GroupMember that is not a team member', :aggregate_failures do
       post :request_access, params: { group_id: group }
 
       expect(controller).to set_flash.to 'Your request for access has been queued for review.'
@@ -469,7 +476,7 @@ RSpec.describe Groups::GroupMembersController do
           group.add_developer(user)
         end
 
-        it 'returns 403' do
+        it 'returns 403', :aggregate_failures do
           post :approve_access_request, params: { group_id: group, id: member }
 
           expect(response).to have_gitlab_http_status(:forbidden)
@@ -482,7 +489,7 @@ RSpec.describe Groups::GroupMembersController do
           group.add_owner(user)
         end
 
-        it 'adds user to members' do
+        it 'adds user to members', :aggregate_failures do
           post :approve_access_request, params: { group_id: group, id: member }
 
           expect(response).to redirect_to(group_group_members_path(group))
diff --git a/spec/controllers/projects/project_members_controller_spec.rb b/spec/controllers/projects/project_members_controller_spec.rb
index bb817fc94b2..9ed43a251a2 100644
--- a/spec/controllers/projects/project_members_controller_spec.rb
+++ b/spec/controllers/projects/project_members_controller_spec.rb
@@ -3,9 +3,9 @@
 require('spec_helper')
 
 RSpec.describe Projects::ProjectMembersController do
-  let(:user) { create(:user) }
-  let(:group) { create(:group, :public) }
-  let(:project) { create(:project, :public) }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:group) { create(:group, :public) }
+  let_it_be(:project, reload: true) { create(:project, :public) }
 
   before do
     travel_to DateTime.new(2019, 4, 1)
@@ -24,8 +24,8 @@ RSpec.describe Projects::ProjectMembersController do
 
     context 'project members' do
       context 'when project belongs to group' do
-        let(:user_in_group) { create(:user) }
-        let(:project_in_group) { create(:project, :public, group: group) }
+        let_it_be(:user_in_group) { create(:user) }
+        let_it_be(:project_in_group) { create(:project, :public, group: group) }
 
         before do
           group.add_owner(user_in_group)
@@ -69,7 +69,7 @@ RSpec.describe Projects::ProjectMembersController do
     end
 
     context 'group links' do
-      let!(:project_group_link) { create(:project_group_link, project: project, group: group) }
+      let_it_be(:project_group_link) { create(:project_group_link, project: project, group: group) }
 
       it 'lists group links' do
         get :index, params: { namespace_id: project.namespace, project_id: project }
@@ -90,7 +90,7 @@ RSpec.describe Projects::ProjectMembersController do
     end
 
     context 'invited members' do
-      let!(:invited_member) { create(:project_member, :invited, project: project) }
+      let_it_be(:invited_member) { create(:project_member, :invited, project: project) }
 
       before do
         project.add_maintainer(user)
@@ -123,7 +123,7 @@ RSpec.describe Projects::ProjectMembersController do
     end
 
     context 'access requests' do
-      let(:access_requester_user) { create(:user) }
+      let_it_be(:access_requester_user) { create(:user) }
 
       before do
         project.request_access(access_requester_user)
@@ -158,7 +158,7 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'POST create' do
-    let(:project_user) { create(:user) }
+    let_it_be(:project_user) { create(:user) }
 
     before do
       sign_in(user)
@@ -169,7 +169,7 @@ RSpec.describe Projects::ProjectMembersController do
         project.add_developer(user)
       end
 
-      it 'returns 404' do
+      it 'returns 404', :aggregate_failures do
         post :create, params: {
                         namespace_id: project.namespace,
                         project_id: project,
@@ -187,11 +187,7 @@ RSpec.describe Projects::ProjectMembersController do
         project.add_maintainer(user)
       end
 
-      it 'adds user to members' do
-        expect_next_instance_of(Members::CreateService) do |instance|
-          expect(instance).to receive(:execute).and_return(status: :success)
-        end
-
+      it 'adds user to members', :aggregate_failures, :snowplow do
         post :create, params: {
                         namespace_id: project.namespace,
                         project_id: project,
@@ -201,9 +197,17 @@ RSpec.describe Projects::ProjectMembersController do
 
         expect(controller).to set_flash.to 'Users were successfully added.'
         expect(response).to redirect_to(project_project_members_path(project))
+        expect(project.users).to include project_user
+        expect_snowplow_event(
+          category: 'Members::CreateService',
+          action: 'create_member',
+          label: 'project-members-page',
+          property: 'existing_user',
+          user: user
+        )
       end
 
-      it 'adds no user to members' do
+      it 'adds no user to members', :aggregate_failures do
         expect_next_instance_of(Members::CreateService) do |instance|
           expect(instance).to receive(:execute).and_return(status: :failure, message: 'Message')
         end
@@ -230,7 +234,7 @@ RSpec.describe Projects::ProjectMembersController do
         unrelated_project.add_maintainer(project_bot)
       end
 
-      it 'returns error' do
+      it 'returns error', :aggregate_failures do
         post :create, params: {
           namespace_id: project.namespace,
           project_id: project,
@@ -261,7 +265,7 @@ RSpec.describe Projects::ProjectMembersController do
       context 'when set to a date in the past' do
         let(:expires_at) { 2.days.ago }
 
-        it 'does not add user to members' do
+        it 'does not add user to members', :aggregate_failures do
           subject
 
           expect(flash[:alert]).to include('Expires at cannot be a date in the past')
@@ -273,7 +277,7 @@ RSpec.describe Projects::ProjectMembersController do
       context 'when set to a date in the future' do
         let(:expires_at) { 5.days.from_now }
 
-        it 'adds user to members' do
+        it 'adds user to members', :aggregate_failures do
           subject
 
           expect(controller).to set_flash.to 'Users were successfully added.'
@@ -285,7 +289,7 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'PUT update' do
-    let(:requester) { create(:project_member, :access_request, project: project) }
+    let_it_be(:requester) { create(:project_member, :access_request, project: project) }
 
     before do
       project.add_maintainer(user)
@@ -393,7 +397,7 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'DELETE destroy' do
-    let(:member) { create(:project_member, :developer, project: project) }
+    let_it_be(:member) { create(:project_member, :developer, project: project) }
 
     before do
       sign_in(user)
@@ -417,7 +421,7 @@ RSpec.describe Projects::ProjectMembersController do
           project.add_developer(user)
         end
 
-        it 'returns 404' do
+        it 'returns 404', :aggregate_failures do
           delete :destroy, params: {
                              namespace_id: project.namespace,
                              project_id: project,
@@ -434,7 +438,7 @@ RSpec.describe Projects::ProjectMembersController do
           project.add_maintainer(user)
         end
 
-        it '[HTML] removes user from members' do
+        it '[HTML] removes user from members', :aggregate_failures do
           delete :destroy, params: {
                              namespace_id: project.namespace,
                              project_id: project,
@@ -447,7 +451,7 @@ RSpec.describe Projects::ProjectMembersController do
           expect(project.members).not_to include member
         end
 
-        it '[JS] removes user from members' do
+        it '[JS] removes user from members', :aggregate_failures do
           delete :destroy, params: {
             namespace_id: project.namespace,
             project_id: project,
@@ -483,7 +487,7 @@ RSpec.describe Projects::ProjectMembersController do
           project.add_developer(user)
         end
 
-        it 'removes user from members' do
+        it 'removes user from members', :aggregate_failures do
           delete :leave, params: {
                            namespace_id: project.namespace,
                            project_id: project
@@ -517,7 +521,7 @@ RSpec.describe Projects::ProjectMembersController do
           project.request_access(user)
         end
 
-        it 'removes user from members' do
+        it 'removes user from members', :aggregate_failures do
           delete :leave, params: {
                            namespace_id: project.namespace,
                            project_id: project
@@ -537,7 +541,7 @@ RSpec.describe Projects::ProjectMembersController do
       sign_in(user)
     end
 
-    it 'creates a new ProjectMember that is not a team member' do
+    it 'creates a new ProjectMember that is not a team member', :aggregate_failures do
       post :request_access, params: {
                               namespace_id: project.namespace,
                               project_id: project
@@ -553,7 +557,7 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'POST approve' do
-    let(:member) { create(:project_member, :access_request, project: project) }
+    let_it_be(:member) { create(:project_member, :access_request, project: project) }
 
     before do
       sign_in(user)
@@ -577,7 +581,7 @@ RSpec.describe Projects::ProjectMembersController do
           project.add_developer(user)
         end
 
-        it 'returns 404' do
+        it 'returns 404', :aggregate_failures do
           post :approve_access_request, params: {
                                           namespace_id: project.namespace,
                                           project_id: project,
@@ -594,7 +598,7 @@ RSpec.describe Projects::ProjectMembersController do
           project.add_maintainer(user)
         end
 
-        it 'adds user to members' do
+        it 'adds user to members', :aggregate_failures do
           post :approve_access_request, params: {
                                           namespace_id: project.namespace,
                                           project_id: project,
@@ -611,8 +615,8 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'POST apply_import' do
-    let(:another_project) { create(:project, :private) }
-    let(:member) { create(:user) }
+    let_it_be(:another_project) { create(:project, :private) }
+    let_it_be(:member) { create(:user) }
 
     before do
       project.add_maintainer(user)
@@ -637,7 +641,7 @@ RSpec.describe Projects::ProjectMembersController do
 
       include_context 'import applied'
 
-      it 'imports source project members' do
+      it 'imports source project members', :aggregate_failures do
         expect(project.team_members).to include member
         expect(controller).to set_flash.to 'Successfully imported'
         expect(response).to redirect_to(
@@ -660,7 +664,7 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'POST create' do
-    let(:stranger) { create(:user) }
+    let_it_be(:stranger) { create(:user) }
 
     context 'when creating owner' do
       before do
@@ -700,7 +704,7 @@ RSpec.describe Projects::ProjectMembersController do
   end
 
   describe 'POST resend_invite' do
-    let(:member) { create(:project_member, project: project) }
+    let_it_be(:member) { create(:project_member, project: project) }
 
     before do
       project.add_maintainer(user)
diff --git a/spec/lib/api/entities/merge_request_basic_spec.rb b/spec/lib/api/entities/merge_request_basic_spec.rb
index 8572b067984..b9d6ab7a652 100644
--- a/spec/lib/api/entities/merge_request_basic_spec.rb
+++ b/spec/lib/api/entities/merge_request_basic_spec.rb
@@ -9,11 +9,22 @@ RSpec.describe ::API::Entities::MergeRequestBasic do
   let_it_be(:labels) { create_list(:label, 3) }
   let_it_be(:merge_requests) { create_list(:labeled_merge_request, 10, :unique_branches, labels: labels) }
 
+  let_it_be(:entity) { described_class.new(merge_request) }
+
   # This mimics the behavior of the `Grape::Entity` serializer
   def present(obj)
     described_class.new(obj).presented
   end
 
+  subject { entity.as_json }
+
+  it 'includes basic fields' do
+    is_expected.to include(
+      draft: merge_request.draft?,
+      work_in_progress: merge_request.draft?
+    )
+  end
+
   context "with :with_api_entity_associations scope" do
     let(:scope) { MergeRequest.with_api_entity_associations }
 
diff --git a/spec/requests/api/merge_requests_spec.rb b/spec/requests/api/merge_requests_spec.rb
index 0c3e74bd1fd..eb6e78aeda2 100644
--- a/spec/requests/api/merge_requests_spec.rb
+++ b/spec/requests/api/merge_requests_spec.rb
@@ -1186,7 +1186,8 @@ RSpec.describe API::MergeRequests do
       expect(json_response['downvotes']).to eq(1)
       expect(json_response['source_project_id']).to eq(merge_request.source_project.id)
       expect(json_response['target_project_id']).to eq(merge_request.target_project.id)
-      expect(json_response['work_in_progress']).to be_falsy
+      expect(json_response['draft']).to be false
+      expect(json_response['work_in_progress']).to be false
       expect(json_response['merge_when_pipeline_succeeds']).to be_falsy
       expect(json_response['merge_status']).to eq('can_be_merged')
       expect(json_response['should_close_merge_request']).to be_falsy
@@ -1329,22 +1330,23 @@ RSpec.describe API::MergeRequests do
       expect(response).to have_gitlab_http_status(:not_found)
     end
 
-    context 'Work in Progress' do
-      let!(:merge_request_wip) do
+    context 'Draft' do
+      let!(:merge_request_draft) do
         create(:merge_request,
           author: user,
           assignees: [user],
           source_project: project,
           target_project: project,
-          title: "WIP: Test",
+          title: "Draft: Test",
           created_at: base_time + 1.second
         )
       end
 
       it "returns merge request" do
-        get api("/projects/#{project.id}/merge_requests/#{merge_request_wip.iid}", user)
+        get api("/projects/#{project.id}/merge_requests/#{merge_request_draft.iid}", user)
 
         expect(response).to have_gitlab_http_status(:ok)
+        expect(json_response['draft']).to eq(true)
         expect(json_response['work_in_progress']).to eq(true)
       end
     end
@@ -2522,8 +2524,8 @@ RSpec.describe API::MergeRequests do
       expect(json_response['message']).to eq('405 Method Not Allowed')
     end
 
-    it "returns 405 if merge_request is a work in progress" do
-      merge_request.update_attribute(:title, "WIP: #{merge_request.title}")
+    it "returns 405 if merge_request is a draft" do
+      merge_request.update_attribute(:title, "Draft: #{merge_request.title}")
       put api("/projects/#{project.id}/merge_requests/#{merge_request.iid}/merge", user)
       expect(response).to have_gitlab_http_status(:method_not_allowed)
       expect(json_response['message']).to eq('405 Method Not Allowed')
diff --git a/spec/services/members/create_service_spec.rb b/spec/services/members/create_service_spec.rb
index 803b0695345..ffe63a8a94b 100644
--- a/spec/services/members/create_service_spec.rb
+++ b/spec/services/members/create_service_spec.rb
@@ -8,7 +8,7 @@ RSpec.describe Members::CreateService, :aggregate_failures, :clean_gitlab_redis_
   let_it_be(:member) { create(:user) }
   let_it_be(:user_ids) { member.id.to_s }
   let_it_be(:access_level) { Gitlab::Access::GUEST }
-  let(:additional_params) { {} }
+  let(:additional_params) { { invite_source: '_invite_source_' } }
   let(:params) { { user_ids: user_ids, access_level: access_level }.merge(additional_params) }
 
   subject(:execute_service) { described_class.new(user, params.merge({ source: source })).execute }
@@ -86,22 +86,16 @@ RSpec.describe Members::CreateService, :aggregate_failures, :clean_gitlab_redis_
 
   context 'when tracking the invite source', :snowplow do
     context 'when invite_source is not passed' do
+      let(:additional_params) { {} }
+
       it 'tracks the invite source as unknown' do
-        execute_service
+        expect { execute_service }.to raise_error(ArgumentError, 'No invite source provided.')
 
-        expect_snowplow_event(
-          category: described_class.name,
-          action: 'create_member',
-          label: 'unknown',
-          property: 'existing_user',
-          user: user
-        )
+        expect_no_snowplow_event
       end
     end
 
-    context 'when invite_source is not passed' do
-      let(:additional_params) { { invite_source: '_invite_source_' } }
-
+    context 'when invite_source is passed' do
       it 'tracks the invite source from params' do
         execute_service
 
@@ -116,7 +110,7 @@ RSpec.describe Members::CreateService, :aggregate_failures, :clean_gitlab_redis_
     end
 
     context 'when it is a net_new_user' do
-      let(:additional_params) { { user_ids: 'email@example.org' } }
+      let(:additional_params) { { invite_source: '_invite_source_', user_ids: 'email@example.org' } }
 
       it 'tracks the invite source from params' do
         execute_service
@@ -124,7 +118,7 @@ RSpec.describe Members::CreateService, :aggregate_failures, :clean_gitlab_redis_
         expect_snowplow_event(
           category: described_class.name,
           action: 'create_member',
-          label: 'unknown',
+          label: '_invite_source_',
           property: 'net_new_user',
           user: user
         )
diff --git a/spec/services/members/invite_service_spec.rb b/spec/services/members/invite_service_spec.rb
index 7007d37ad8d..c530e3d0c53 100644
--- a/spec/services/members/invite_service_spec.rb
+++ b/spec/services/members/invite_service_spec.rb
@@ -8,7 +8,7 @@ RSpec.describe Members::InviteService, :aggregate_failures, :clean_gitlab_redis_
   let_it_be(:project_user) { create(:user) }
   let_it_be(:namespace) { project.namespace }
   let(:params) { {} }
-  let(:base_params) { { access_level: Gitlab::Access::GUEST, source: project } }
+  let(:base_params) { { access_level: Gitlab::Access::GUEST, source: project, invite_source: '_invite_source_' } }
 
   subject(:result) { described_class.new(user, base_params.merge(params) ).execute }
 
diff --git a/spec/tasks/gitlab/storage_rake_spec.rb b/spec/tasks/gitlab/storage_rake_spec.rb
index 9dee27ca391..fefcd05af3b 100644
--- a/spec/tasks/gitlab/storage_rake_spec.rb
+++ b/spec/tasks/gitlab/storage_rake_spec.rb
@@ -88,6 +88,27 @@ RSpec.describe 'rake gitlab:storage:*', :silence_stdout 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:*', :silence_stdout 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:*', :silence_stdout 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