1 files changed, 12 insertions, 6 deletions
diff --git a/spec/frontend/notes/components/discussion_notes_spec.js b/spec/frontend/notes/components/discussion_notes_spec.js
index 51c7f45d5b0..d259e79de84 100644
--- a/spec/frontend/notes/components/discussion_notes_spec.js
+++ b/spec/frontend/notes/components/discussion_notes_spec.js
@@ -6,7 +6,6 @@ import NoteableNote from '~/notes/components/noteable_note.vue';
 import PlaceholderNote from '~/vue_shared/components/notes/placeholder_note.vue';
 import PlaceholderSystemNote from '~/vue_shared/components/notes/placeholder_system_note.vue';
 import SystemNote from '~/vue_shared/components/notes/system_note.vue';
-import TimelineEntryItem from '~/vue_shared/components/notes/timeline_entry_item.vue';
 import createStore from '~/notes/stores';
 import { noteableDataMock, discussionMock, notesDataMock } from '../../notes/mock_data';
 
@@ -48,13 +47,13 @@ describe('DiscussionNotes', () => {
     it('renders an element for each note in the discussion', () => {
       createComponent();
       const notesCount = discussionMock.notes.length;
-      const els = wrapper.findAll(TimelineEntryItem);
+      const els = wrapper.findAll(NoteableNote);
       expect(els.length).toBe(notesCount);
     });
 
     it('renders one element if replies groupping is enabled', () => {
       createComponent({ shouldGroupReplies: true });
-      const els = wrapper.findAll(TimelineEntryItem);
+      const els = wrapper.findAll(NoteableNote);
       expect(els.length).toBe(1);
     });
 
@@ -85,7 +84,7 @@ describe('DiscussionNotes', () => {
       ];
       discussion.notes = notesData;
       createComponent({ discussion, shouldRenderDiffs: true });
-      const notes = wrapper.findAll('.notes > li');
+      const notes = wrapper.findAll('.notes > *');
 
       expect(notes.at(0).is(PlaceholderSystemNote)).toBe(true);
       expect(notes.at(1).is(PlaceholderNote)).toBe(true);
@@ -111,7 +110,14 @@ describe('DiscussionNotes', () => {
 
   describe('events', () => {
     describe('with groupped notes and replies expanded', () => {
-      const findNoteAtIndex = index => wrapper.find(`.note:nth-of-type(${index + 1}`);
+      const findNoteAtIndex = index => {
+        const noteComponents = [NoteableNote, SystemNote, PlaceholderNote, PlaceholderSystemNote];
+        const allowedNames = noteComponents.map(c => c.name);
+        return wrapper
+          .findAll('.notes *')
+          .filter(w => allowedNames.includes(w.name()))
+          .at(index);
+      };
 
       beforeEach(() => {
         createComponent({ shouldGroupReplies: true, isExpanded: true });
@@ -146,7 +152,7 @@ describe('DiscussionNotes', () => {
       let note;
       beforeEach(() => {
         createComponent();
-        note = wrapper.find('.note');
+        note = wrapper.find('.notes > *');
       });
 
       it('emits deleteNote when first note emits handleDeleteNote', () => {