round of fixes after code review

6 files changed, 61 insertions, 26 deletions

diff --git a/app/controllers/concerns/renders_notes.rb b/app/controllers/concerns/renders_notes.rb
index cdcfefd90c9..b6cf366c05c 100644
--- a/app/controllers/concerns/renders_notes.rb
+++ b/app/controllers/concerns/renders_notes.rb
@@ -1,5 +1,5 @@
 module RendersNotes
-  def prepare_notes_for_rendering(notes, noteable=nil)
+  def prepare_notes_for_rendering(notes, noteable = nil)
     preload_noteable_for_regular_notes(notes)
     preload_max_access_for_authors(notes, @project)
     preload_first_time_contribution_for_authors(noteable, notes) if noteable.is_a?(Issuable)
@@ -23,7 +23,10 @@ module RendersNotes
 
   def preload_first_time_contribution_for_authors(issuable, notes)
     return unless issuable.first_contribution?
+
     same_author = lambda {|n| n.author_id == issuable.author_id}
-    notes.select(&same_author).each {|note| note.special_role = :first_time_contributor}
+    notes.each do |note|
+      note.specialize!(Note::SpecialRole::FIRST_TIME_CONTRIBUTOR, &same_author)
+    end
   end
 end
diff --git a/app/models/concerns/issuable.rb b/app/models/concerns/issuable.rb
index 7b1c4e7a2d5..0cc60dddc69 100644
--- a/app/models/concerns/issuable.rb
+++ b/app/models/concerns/issuable.rb
@@ -337,6 +337,7 @@ module Issuable
 
   def first_contribution?
     return false if project.team.max_member_access(author_id) > Gitlab::Access::GUEST
+
     project.merge_requests.merged.where(author_id: author_id).empty?
   end
 end
diff --git a/app/models/note.rb b/app/models/note.rb
index d04bccc6e81..28cd54bd81c 100644
--- a/app/models/note.rb
+++ b/app/models/note.rb
@@ -26,7 +26,6 @@ class Note < ActiveRecord::Base
   end
 
   ignore_column :original_discussion_id
-  ignore_column :special_role
 
   cache_markdown_field :note, pipeline: :note, issuable_state_filter_enabled: true
 
@@ -155,6 +154,10 @@ class Note < ActiveRecord::Base
         .group(:noteable_id)
         .where(noteable_type: type, noteable_id: ids)
     end
+
+    def has_special_role?(role, note)
+      note.special_role == role
+    end
   end
 
   def cross_reference?
@@ -221,14 +224,19 @@ class Note < ActiveRecord::Base
   end
 
   def special_role=(role)
-    raise "Role is undefined, #{role} not found in #{SpecialRole.values}" unless SpecialRole.values.include? role
+    raise "Role is undefined, #{role} not found in #{SpecialRole.values}" unless SpecialRole.values.include?(role)
+
     @special_role = role
   end
 
   def has_special_role?(role)
-    return @special_role == role
+    self.class.has_special_role?(role, self)
   end
 
+  def specialize!(role)
+    self.special_role = role if !block_given? || yield(self)
+  end
+ 
   def editable?
     !system?
   end
diff --git a/app/views/projects/notes/_actions.html.haml b/app/views/projects/notes/_actions.html.haml
index 926ceff3ee2..cf634faccfc 100644
--- a/app/views/projects/notes/_actions.html.haml
+++ b/app/views/projects/notes/_actions.html.haml
@@ -1,7 +1,7 @@
-- if note.has_special_role? :first_time_contributor
+- if note.has_special_role?(Note::SpecialRole::FIRST_TIME_CONTRIBUTOR)
   %span.note-role.has-tooltip{ title: _("This user hasn't yet contributed code to this project. Handle with care.")}
     = _('First-time contributor')
-- elsif access = access = note_max_access_for_user(note)
+- elsif access = note_max_access_for_user(note)
   %span.note-role= Gitlab::Access.human_access(access)
 
 - if note.resolvable?
diff --git a/spec/models/concerns/issuable_spec.rb b/spec/models/concerns/issuable_spec.rb
index 9b83b8bfb7e..9948340d976 100644
--- a/spec/models/concerns/issuable_spec.rb
+++ b/spec/models/concerns/issuable_spec.rb
@@ -492,15 +492,14 @@ describe Issuable do
 
     let(:contributor) { create(:user) }
     let(:first_time_contributor) { create(:user) }
-    let!(:access_users) { [owner, master, reporter] }
 
     before do
       group.add_owner(owner)
-      project.team << [master, :master]
-      project.team << [reporter, :reporter]
-      project.team << [guest, :guest]
-      project.team << [contributor, :guest]
-      project.team << [first_time_contributor, :guest]
+      project.add_master(master)
+      project.add_reporter(reporter)
+      project.add_guest(guest)
+      project.add_guest(contributor)
+      project.add_guest(first_time_contributor)
     end
     
     let(:merged_mr) { create(:merge_request, :merged, author: contributor, target_project: project, source_project: project) }
@@ -510,27 +509,30 @@ describe Issuable do
     context "for merge requests" do
       it "is false for MASTER" do
         mr = create(:merge_request, author: master, target_project: project, source_project: project)
-        expect(mr.first_contribution?).to be_falsey
+
+        expect(mr).not_to be_first_contribution
       end
 
       it "is false for OWNER" do
         mr = create(:merge_request, author: owner, target_project: project, source_project: project)
-        expect(mr.first_contribution?).to be_falsey
+        
+        expect(mr).not_to be_first_contribution
       end
 
       it "is false for REPORTER" do
         mr = create(:merge_request, author: reporter, target_project: project, source_project: project)
-        expect(mr.first_contribution?).to be_falsey
+        
+        expect(mr).not_to be_first_contribution
       end
 
       it "is true when you don't have any merged MR" do
-        expect(open_mr.first_contribution?).to be_truthy
-        expect(merged_mr.first_contribution?).to be_falsey
+        expect(open_mr).to be_first_contribution
+        expect(merged_mr).not_to be_first_contribution
       end
 
-      it "handle multiple projects separately" do
-        expect(open_mr.first_contribution?).to be_truthy
-        expect(merged_mr_other_project.first_contribution?).to be_falsey
+      it "handles multiple projects separately" do
+        expect(open_mr).to be_first_contribution
+        expect(merged_mr_other_project).not_to be_first_contribution
       end
     end
 
@@ -541,15 +543,15 @@ describe Issuable do
 
       it "is true when you don't have any merged MR" do
         expect(merged_mr).to be
-        expect(first_time_contributor_issue.first_contribution?).to be_truthy
-        expect(contributor_issue.first_contribution?).to be_falsey
+        expect(first_time_contributor_issue).to be_first_contribution
+        expect(contributor_issue).not_to be_first_contribution
       end
 
-      it "handle multiple projects separately" do
+      it "handles multiple projects separately" do
         expect(merged_mr).to be
         expect(merged_mr_other_project).to be
-        expect(first_time_contributor_issue.first_contribution?).to be_truthy
-        expect(first_time_contributor_issue_other_project.first_contribution?).to be_falsey
+        expect(first_time_contributor_issue).to be_first_contribution
+        expect(first_time_contributor_issue_other_project).not_to be_first_contribution
       end
     end
   end
diff --git a/spec/models/note_spec.rb b/spec/models/note_spec.rb
index b214074fdce..09caa2c9991 100644
--- a/spec/models/note_spec.rb
+++ b/spec/models/note_spec.rb
@@ -696,4 +696,25 @@ describe Note do
       note.destroy!
     end
   end
+
+  describe '#specialize!' do
+    let(:note) { build(:note) }
+    let(:role) { Note::SpecialRole::FIRST_TIME_CONTRIBUTOR }
+
+    it 'should set special role without a block' do
+      expect { note.specialize!(role) }.to change { note.special_role }.from(nil).to(role)
+    end
+
+    it 'should set special role with a truthy block' do
+      tautology = -> (*) { true }
+      
+      expect { note.specialize!(role, &tautology) }.to change { note.special_role }.from(nil).to(role)
+    end
+
+    it 'should not set special role with a falsey block' do
+      contradiction = -> (*) { false }
+
+      expect { note.specialize!(role, &contradiction) }.not_to change { note.special_role }
+    end
+  end
 end