1 files changed, 6 insertions, 3 deletions
diff --git a/spec/frontend/image_diff/helpers/badge_helper_spec.js b/spec/frontend/image_diff/helpers/badge_helper_spec.js
index c970ccc535d..9ac6ebf6fdb 100644
--- a/spec/frontend/image_diff/helpers/badge_helper_spec.js
+++ b/spec/frontend/image_diff/helpers/badge_helper_spec.js
@@ -62,7 +62,10 @@ describe('badge helper', () => {
     });
 
     it('should add badge classes', () => {
-      expect(buttonEl.className).toContain('badge badge-pill');
+      const classes = buttonEl.className.split(' ');
+      expect(classes).toEqual(
+        expect.arrayContaining(['design-note-pin', 'on-image', 'gl-absolute']),
+      );
     });
 
     it('should set the badge text', () => {
@@ -105,7 +108,7 @@ describe('badge helper', () => {
     beforeEach(() => {
       containerEl.innerHTML = `
         <div id="${noteId}">
-          <div class="badge hidden">
+          <div class="design-note-pin hidden">
           </div>
         </div>
       `;
@@ -116,7 +119,7 @@ describe('badge helper', () => {
           badgeNumber,
         },
       });
-      avatarBadgeEl = containerEl.querySelector(`#${noteId} .badge`);
+      avatarBadgeEl = containerEl.querySelector(`#${noteId} .design-note-pin`);
     });
 
     it('should update badge number', () => {