1 files changed, 14 insertions, 8 deletions
diff --git a/spec/features/merge_requests/diff_notes_avatars_spec.rb b/spec/features/merge_requests/diff_notes_avatars_spec.rb
index 9bcb78d5206..4766cdf716f 100644
--- a/spec/features/merge_requests/diff_notes_avatars_spec.rb
+++ b/spec/features/merge_requests/diff_notes_avatars_spec.rb
@@ -84,7 +84,7 @@ feature 'Diff note avatars', js: true do
       end
 
       it 'shows note avatar' do
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           find('.diff-notes-collapse').click
 
           expect(page).to have_selector('img.js-diff-comment-avatar', count: 1)
@@ -92,7 +92,7 @@ feature 'Diff note avatars', js: true do
       end
 
       it 'shows comment on note avatar' do
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           find('.diff-notes-collapse').click
 
           expect(first('img.js-diff-comment-avatar')["data-original-title"]).to eq("#{note.author.name}: #{note.note.truncate(17)}")
@@ -100,13 +100,13 @@ feature 'Diff note avatars', js: true do
       end
 
       it 'toggles comments when clicking avatar' do
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           find('.diff-notes-collapse').click
         end
 
         expect(page).to have_selector('.notes_holder', visible: false)
 
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           first('img.js-diff-comment-avatar').click
         end
 
@@ -122,7 +122,7 @@ feature 'Diff note avatars', js: true do
 
         wait_for_requests
 
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           expect(page).not_to have_selector('img.js-diff-comment-avatar')
         end
       end
@@ -138,7 +138,7 @@ feature 'Diff note avatars', js: true do
           wait_for_requests
         end
 
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           find('.diff-notes-collapse').trigger('click')
 
           expect(page).to have_selector('img.js-diff-comment-avatar', count: 2)
@@ -158,7 +158,7 @@ feature 'Diff note avatars', js: true do
           end
         end
 
-        page.within find("[id='#{position.line_code(project.repository)}']") do
+        page.within find_line(position.line_code(project.repository)) do
           find('.diff-notes-collapse').trigger('click')
 
           expect(page).to have_selector('img.js-diff-comment-avatar', count: 3)
@@ -176,7 +176,7 @@ feature 'Diff note avatars', js: true do
         end
 
         it 'shows extra comment count' do
-          page.within find("[id='#{position.line_code(project.repository)}']") do
+          page.within find_line(position.line_code(project.repository)) do
             find('.diff-notes-collapse').click
 
             expect(find('.diff-comments-more-count')).to have_content '+1'
@@ -185,4 +185,10 @@ feature 'Diff note avatars', js: true do
       end
     end
   end
+
+  def find_line(line_code)
+    line = find("[id='#{line_code}']")
+    line = line.find(:xpath, 'preceding-sibling::*[1][self::td]') if line.tag_name == 'td'
+    line
+  end
 end