1 files changed, 20 insertions, 9 deletions
diff --git a/spec/models/merge_request_spec.rb b/spec/models/merge_request_spec.rb
index 98a39c33319..bb63abd167b 100644
--- a/spec/models/merge_request_spec.rb
+++ b/spec/models/merge_request_spec.rb
@@ -601,30 +601,30 @@ describe MergeRequest do
   end
 
   describe '#can_remove_source_branch?' do
-    let(:user) { create(:user) }
-    let(:user2) { create(:user) }
+    set(:user) { create(:user) }
+    set(:merge_request) { create(:merge_request, :simple) }
 
-    before do
-      subject.source_project.team << [user, :master]
+    subject { merge_request }
 
-      subject.source_branch = "feature"
-      subject.target_branch = "master"
-      subject.save!
+    before do
+      subject.source_project.add_master(user)
     end
 
     it "can't be removed when its a protected branch" do
       allow(ProtectedBranch).to receive(:protected?).and_return(true)
+
       expect(subject.can_remove_source_branch?(user)).to be_falsey
     end
 
     it "can't remove a root ref" do
-      subject.source_branch = "master"
-      subject.target_branch = "feature"
+      subject.update(source_branch: 'master', target_branch: 'feature')
 
       expect(subject.can_remove_source_branch?(user)).to be_falsey
     end
 
     it "is unable to remove the source branch for a project the user cannot push to" do
+      user2 = create(:user)
+
       expect(subject.can_remove_source_branch?(user2)).to be_falsey
     end
 
@@ -635,6 +635,7 @@ describe MergeRequest do
     end
 
     it "cannot be removed if the last commit is not also the head of the source branch" do
+      subject.clear_memoized_shas
       subject.source_branch = "lfs"
 
       expect(subject.can_remove_source_branch?(user)).to be_falsey
@@ -1405,6 +1406,16 @@ describe MergeRequest do
 
       subject.reload_diff
     end
+
+    context 'when using the after_update hook to update' do
+      context 'when the branches are updated' do
+        it 'uses the new heads to generate the diff' do
+          expect { subject.update!(source_branch: subject.target_branch, target_branch: subject.source_branch) }
+            .to change { subject.merge_request_diff.start_commit_sha }
+            .and change { subject.merge_request_diff.head_commit_sha }
+        end
+      end
+    end
   end
 
   describe '#update_diff_discussion_positions' do