diff --git a/src/renderer/components/update-button/__tests__/__snapshots__/update-button.test.tsx.snap b/src/renderer/components/update-button/__tests__/__snapshots__/update-button.test.tsx.snap
index 6335f53d46..3492065df4 100644
--- a/src/renderer/components/update-button/__tests__/__snapshots__/update-button.test.tsx.snap
+++ b/src/renderer/components/update-button/__tests__/__snapshots__/update-button.test.tsx.snap
@@ -4,6 +4,7 @@ exports[`<UpdateButton/> should render if warning level prop passed 1`] = `
 <button
   class="updateButton"
   data-testid="update-button"
+  data-warning-level="light"
   id="update-lens-button"
 >
   Update
diff --git a/src/renderer/components/update-button/__tests__/update-button.test.tsx b/src/renderer/components/update-button/__tests__/update-button.test.tsx
index d476ab09ef..98a86aa318 100644
--- a/src/renderer/components/update-button/__tests__/update-button.test.tsx
+++ b/src/renderer/components/update-button/__tests__/update-button.test.tsx
@@ -51,19 +51,27 @@ describe("<UpdateButton/>", () => {
     expect(update).toHaveBeenCalled();
   });
 
-  it("should have class name with medium warning level", () => {
+  it("should have light warning level", () => {
+    const { getByTestId } = render(<UpdateButton update={update} warningLevel="light" />);
+
+    const button = getByTestId("update-button");
+
+    expect(button.dataset.warningLevel).toBe("light");
+  });
+
+  it("should have medium warning level", () => {
     const { getByTestId } = render(<UpdateButton update={update} warningLevel="medium" />);
 
     const button = getByTestId("update-button");
 
-    expect(button.className.toLowerCase().includes("medium")).toBeTruthy();
+    expect(button.dataset.warningLevel).toBe("medium");
   });
 
-  it("should have class name with high warning level", () => {
+  it("should have high warning level", () => {
     const { getByTestId } = render(<UpdateButton update={update} warningLevel="high" />);
 
     const button = getByTestId("update-button");
 
-    expect(button.className.toLowerCase().includes("high")).toBeTruthy();
+    expect(button.dataset.warningLevel).toBe("high");
   });
 });
diff --git a/src/renderer/components/update-button/update-button.tsx b/src/renderer/components/update-button/update-button.tsx
index 74313992a7..6da453de3d 100644
--- a/src/renderer/components/update-button/update-button.tsx
+++ b/src/renderer/components/update-button/update-button.tsx
@@ -34,6 +34,7 @@ export function UpdateButton({ warningLevel, update, id }: UpdateButtonProps) {
     <>
       <button
         data-testid="update-button"
+        data-warning-level={warningLevel}
         id={buttonId}
         className={cssNames(styles.updateButton, {
           [styles.warningHigh]: warningLevel === "high",