diff --git a/code/ui/blocks/src/blocks/internal/InternalCanvas.stories.tsx b/code/ui/blocks/src/blocks/internal/InternalCanvas.stories.tsx
index dc9a67ff79ac..51107dc87bcf 100644
--- a/code/ui/blocks/src/blocks/internal/InternalCanvas.stories.tsx
+++ b/code/ui/blocks/src/blocks/internal/InternalCanvas.stories.tsx
@@ -2,7 +2,7 @@
 /// <reference types="@testing-library/jest-dom" />;
 import React from 'react';
 import type { Meta, StoryObj } from '@storybook/react';
-import { userEvent, waitFor, within } from '@storybook/testing-library';
+import { userEvent, within } from '@storybook/testing-library';
 import { expect } from '@storybook/jest';
 import { Canvas } from '../Canvas';
 import { Story as StoryComponent } from '../Story';
@@ -12,6 +12,7 @@ const meta: Meta<typeof Canvas> = {
   title: 'Blocks/Internal/Canvas',
   component: Canvas,
   parameters: {
+    theme: 'light',
     relativeCsfPaths: ['../examples/Button.stories'],
   },
   render: (args) => {
@@ -32,20 +33,15 @@ const expectAmountOfStoriesInSource =
     const canvas = within(canvasElement);
 
     // Arrange - find the "Show code" button
-    let showCodeButton = canvas.getByText('Show code');
-    await waitFor(() => {
-      showCodeButton = canvas.getByText('Show code');
-      expect(showCodeButton).toBeInTheDocument();
-    });
+    const showCodeButton = await canvas.findByText('Show code');
+    await expect(showCodeButton).toBeInTheDocument();
 
     // Act - click button to show code
     await userEvent.click(showCodeButton);
 
     // Assert - check that the correct amount of stories' source is shown
-    await waitFor(async () => {
-      const booleanControlNodes = await canvas.findAllByText('BooleanControl');
-      expect(booleanControlNodes).toHaveLength(amount);
-    });
+    const booleanControlNodes = await canvas.findAllByText('onClick');
+    await expect(booleanControlNodes).toHaveLength(amount);
   };
 
 export const MultipleChildren: Story = {
@@ -115,6 +111,6 @@ export const MixedChildrenStories: Story = {
     const canvas = within(args.canvasElement);
 
     // Assert - only find two headlines, those in the story, and none in the source code
-    expect(canvas.queryAllByText(/Headline for Boolean Controls/i)).toHaveLength(2);
+    expect(canvas.queryAllByText(/Headline for /i)).toHaveLength(2);
   },
 };