1 files changed, 8 insertions, 4 deletions
diff --git a/spec/frontend/image_diff/helpers/dom_helper_spec.js b/spec/frontend/image_diff/helpers/dom_helper_spec.js
index 9357d626bbe..1c5f1cbe3da 100644
--- a/spec/frontend/image_diff/helpers/dom_helper_spec.js
+++ b/spec/frontend/image_diff/helpers/dom_helper_spec.js
@@ -37,14 +37,16 @@ describe('domHelper', () => {
       discussionEl = document.createElement('div');
       discussionEl.innerHTML = `
         <a href="#" class="image-diff-avatar-link">
-          <div class="badge"></div>
+          <div class="design-note-pin"></div>
         </a>
       `;
       domHelper.updateDiscussionAvatarBadgeNumber(discussionEl, badgeNumber);
     });
 
     it('should update avatar badge number', () => {
-      expect(discussionEl.querySelector('.badge').textContent).toEqual(badgeNumber.toString());
+      expect(discussionEl.querySelector('.design-note-pin').textContent).toEqual(
+        badgeNumber.toString(),
+      );
     });
   });
 
@@ -54,13 +56,15 @@ describe('domHelper', () => {
     beforeEach(() => {
       discussionEl = document.createElement('div');
       discussionEl.innerHTML = `
-        <div class="badge"></div>
+        <div class="design-note-pin"></div>
       `;
       domHelper.updateDiscussionBadgeNumber(discussionEl, badgeNumber);
     });
 
     it('should update discussion badge number', () => {
-      expect(discussionEl.querySelector('.badge').textContent).toEqual(badgeNumber.toString());
+      expect(discussionEl.querySelector('.design-note-pin').textContent).toEqual(
+        badgeNumber.toString(),
+      );
     });
   });