diff --git a/.storybook/test-runner.ts b/.storybook/test-runner.ts
index f291980f..ca9f5bea 100644
--- a/.storybook/test-runner.ts
+++ b/.storybook/test-runner.ts
@@ -1,14 +1,15 @@
 import { toMatchImageSnapshot } from 'jest-image-snapshot';
 import type { TestRunnerConfig } from '../dist/ts';
 
-const customSnapshotsDir = `${process.cwd()}/__snapshots__`;
+const snapshotsDir = process.env.SNAPSHOTS_DIR || '__snapshots__';
+const customSnapshotsDir = `${process.cwd()}/${snapshotsDir}`;
 
 const config: TestRunnerConfig = {
   setup() {
     expect.extend({ toMatchImageSnapshot });
   },
   async postRender(page, context) {
-    const image = await page.screenshot();
+    const image = await page.screenshot({ fullPage: true });
     expect(image).toMatchImageSnapshot({
       customSnapshotsDir,
       customSnapshotIdentifier: context.id,
diff --git a/src/setup-page.ts b/src/setup-page.ts
index 87d6ac7e..f6b27b82 100644
--- a/src/setup-page.ts
+++ b/src/setup-page.ts
@@ -20,6 +20,8 @@ const sanitizeURL = (url) => {
 export const setupPage = async (page) => {
   const start = new Date();
   const targetURL = sanitizeURL(process.env.TARGET_URL || `http://localhost:6006`);
+  const viewMode = process.env.VIEW_MODE || 'story';
+  const renderedEvent = viewMode === 'docs' ? 'docsRendered' : 'storyRendered';
 
   const referenceURL = process.env.REFERENCE_URL && sanitizeURL(process.env.REFERENCE_URL);
 
@@ -103,7 +105,7 @@ export const setupPage = async (page) => {
         }
 
         return new Promise((resolve, reject) => {
-          channel.on('storyRendered', () => resolve(document.getElementById('root')));
+          channel.on('${renderedEvent}', () => resolve(document.getElementById('root')));
           channel.on('storyUnchanged', () => resolve(document.getElementById('root')));
           channel.on('storyErrored', ({ description }) => reject(
             new StorybookTestRunnerError(storyId, description))
@@ -115,7 +117,7 @@ export const setupPage = async (page) => {
             new StorybookTestRunnerError(storyId, 'The story was missing when trying to access it.'))
           );
 
-          channel.emit('setCurrentStory', { storyId });
+          channel.emit('setCurrentStory', { storyId, viewMode: '${viewMode}' });
         });
       };
     `,