1 files changed, 10 insertions, 9 deletions
diff --git a/spec/frontend/batch_comments/components/diff_file_drafts_spec.js b/spec/frontend/batch_comments/components/diff_file_drafts_spec.js
index dcb68b1804f..6a5ff1af7c9 100644
--- a/spec/frontend/batch_comments/components/diff_file_drafts_spec.js
+++ b/spec/frontend/batch_comments/components/diff_file_drafts_spec.js
@@ -1,11 +1,11 @@
-import { shallowMount, createLocalVue } from '@vue/test-utils';
+import { shallowMount } from '@vue/test-utils';
+import Vue from 'vue';
 import Vuex from 'vuex';
 import DiffFileDrafts from '~/batch_comments/components/diff_file_drafts.vue';
 import DraftNote from '~/batch_comments/components/draft_note.vue';
+import DesignNotePin from '~/vue_shared/components/design_management/design_note_pin.vue';
 
-const localVue = createLocalVue();
-
-localVue.use(Vuex);
+Vue.use(Vuex);
 
 describe('Batch comments diff file drafts component', () => {
   let vm;
@@ -22,9 +22,8 @@ describe('Batch comments diff file drafts component', () => {
       },
     });
 
-    vm = shallowMount(localVue.extend(DiffFileDrafts), {
+    vm = shallowMount(DiffFileDrafts, {
       store,
-      localVue,
       propsData: { fileHash: 'filehash' },
     });
   }
@@ -42,10 +41,12 @@ describe('Batch comments diff file drafts component', () => {
   it('renders index of draft note', () => {
     factory();
 
-    expect(vm.findAll('.js-diff-notes-index').length).toEqual(2);
+    const elements = vm.findAll(DesignNotePin);
+
+    expect(elements.length).toEqual(2);
 
-    expect(vm.findAll('.js-diff-notes-index').at(0).text()).toEqual('1');
+    expect(elements.at(0).props('label')).toEqual(1);
 
-    expect(vm.findAll('.js-diff-notes-index').at(1).text()).toEqual('2');
+    expect(elements.at(1).props('label')).toEqual(2);
   });
 });