Add more tests to merge_request_diff and improve initialize

app/models/merge_request.rb

@@ -246,7 +246,9 @@ class MergeRequest < ActiveRecord::Base
   end
 
   def diff_refs
-    merge_request_diff.diff_refs
+    if merge_request_diff
+      merge_request_diff.diff_refs
+    end
   end
 
   def validate_branches

app/models/merge_request_diff.rb

@@ -35,7 +35,7 @@ class MergeRequestDiff < ActiveRecord::Base
   # have those variables in the database so we try to initialize it
   def initialize_commits_sha
     self.start_commit_sha ||= merge_request.target_branch_sha
-    self.head_commit_sha  ||= last_commit.try(:sha) || merge_request.source_branch_sha
+    self.head_commit_sha  ||= persisted? ? last_commit.sha : merge_request.source_branch_sha
     self.base_commit_sha  ||= find_base_sha
   end
 
@@ -191,9 +191,7 @@ class MergeRequestDiff < ActiveRecord::Base
     end
 
     new_attributes[:st_diffs] = new_diffs
-
     update_columns_serialized(new_attributes)
-
     keep_around_commits
   end
 

spec/models/merge_request_diff_spec.rb

 require 'spec_helper'
 
 describe MergeRequestDiff, models: true do
+  describe 'initialize new object' do
+    subject { build(:merge_request).merge_request_diffs.build }
+
+    it { expect(subject).to be_valid }
+    it { expect(subject.head_commit_sha).to eq('5937ac0a7beb003549fc5fd26fc247adbce4a52e') }
+    it { expect(subject.base_commit_sha).to eq('ae73cb07c9eeaf35924a10f713b364d32b2dd34f') }
+    it { expect(subject.start_commit_sha).to eq('0b4bc9a49b562e85de7cc9e834518ea6828729b9') }
+  end
+
+  describe 'create new record' do
+    subject { create(:merge_request) }
+
+    it { expect(subject).to be_valid }
+    it { expect(subject.commits.count).to eq(5) }
+    it { expect(subject.diffs.count).to eq(8) }
+  end
+
   describe '#diffs' do
     let(:mr) { create(:merge_request, :with_diffs) }
     let(:mr_diff) { mr.merge_request_diff }