diff --git a/src/internal/visual-mode/__tests__/use-runtime-visual-refresh.test.tsx b/src/internal/visual-mode/__tests__/use-runtime-visual-refresh.test.tsx
index b6e4bfa..1a56b5b 100644
--- a/src/internal/visual-mode/__tests__/use-runtime-visual-refresh.test.tsx
+++ b/src/internal/visual-mode/__tests__/use-runtime-visual-refresh.test.tsx
@@ -18,8 +18,10 @@ describe('useVisualRefresh', () => {
     return <div data-testid="current-mode">{isRefresh.toString()}</div>;
   }
 
-  beforeEach(() => clearVisualRefreshState());
-  afterEach(() => document.body.classList.remove('awsui-visual-refresh'));
+  afterEach(() => {
+    clearVisualRefreshState();
+    expect(document.querySelector('.awsui-visual-refresh')).toBeFalsy();
+  });
   afterEach(() => {
     clearMessageCache();
     jest.restoreAllMocks();
diff --git a/src/internal/visual-mode/index.ts b/src/internal/visual-mode/index.ts
index 81aff69..468116b 100644
--- a/src/internal/visual-mode/index.ts
+++ b/src/internal/visual-mode/index.ts
@@ -107,6 +107,9 @@ let visualRefreshState: undefined | boolean = undefined;
 // for testing
 export function clearVisualRefreshState() {
   visualRefreshState = undefined;
+  if (typeof document !== 'undefined') {
+    document.body.classList.remove('awsui-visual-refresh');
+  }
 }
 
 function detectVisualRefreshClassName() {