From 1385c6d3ed7c6839d855463a223e21e2cea6b8f5 Mon Sep 17 00:00:00 2001
From: Yann Braga <yannbf@gmail.com>
Date: Wed, 24 Aug 2022 21:59:19 +0200
Subject: [PATCH] support storystorev7 in coverage check

---
 src/playwright/transformPlaywright.test.ts    |  24 +
 src/playwright/transformPlaywright.ts         |  12 +
 .../transformPlaywrightJson.test.ts           | 492 ++++++++++--------
 src/setup-page.ts                             |  13 -
 4 files changed, 314 insertions(+), 227 deletions(-)

diff --git a/src/playwright/transformPlaywright.test.ts b/src/playwright/transformPlaywright.test.ts
index 75c1d7fb..63967496 100644
--- a/src/playwright/transformPlaywright.test.ts
+++ b/src/playwright/transformPlaywright.test.ts
@@ -89,6 +89,14 @@ describe('Playwright', () => {
                 }
         
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+        
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+        
                   await jestPlaywright.saveCoverage(page);
                 }
         
@@ -164,6 +172,14 @@ describe('Playwright', () => {
                 }
         
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+        
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+        
                   await jestPlaywright.saveCoverage(page);
                 }
         
@@ -240,6 +256,14 @@ describe('Playwright', () => {
                 }
         
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+        
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+        
                   await jestPlaywright.saveCoverage(page);
                 }
         
diff --git a/src/playwright/transformPlaywright.ts b/src/playwright/transformPlaywright.ts
index c04f6537..b6ee8023 100644
--- a/src/playwright/transformPlaywright.ts
+++ b/src/playwright/transformPlaywright.ts
@@ -4,12 +4,19 @@ import { userOrAutoTitle } from '@storybook/store';
 
 import { getStorybookMetadata } from '../util';
 import { transformCsf } from '../csf/transformCsf';
+import dedent from 'ts-dedent';
 
 const filePrefixer = template(`
   import global from 'global';
   const { setupPage } = require('@storybook/test-runner');
 `);
 
+const coverageErrorMessage = dedent`
+  [Test runner] An error occurred when evaluating code coverage:
+  The code in this story is not instrumented, which means the coverage setup is likely not correct.
+  More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage
+`;
+
 export const testPrefixer = template(
   `
     console.log({ id: %%id%%, title: %%title%%, name: %%name%%, storyExport: %%storyExport%% });
@@ -34,6 +41,11 @@ export const testPrefixer = template(
         }
 
         if(global.__sbCollectCoverage) {
+          const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+          if (!isCoverageSetupCorrectly) {
+            throw new Error(\`${coverageErrorMessage}\`);
+          }
+
           await jestPlaywright.saveCoverage(page);
         }
 
diff --git a/src/playwright/transformPlaywrightJson.test.ts b/src/playwright/transformPlaywrightJson.test.ts
index c846203f..bbabd6ae 100644
--- a/src/playwright/transformPlaywrightJson.test.ts
+++ b/src/playwright/transformPlaywrightJson.test.ts
@@ -63,6 +63,14 @@ describe('Playwright Json', () => {
                 }
 
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
                   await jestPlaywright.saveCoverage(page);
                 }
 
@@ -116,6 +124,14 @@ describe('Playwright Json', () => {
                 }
 
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
                   await jestPlaywright.saveCoverage(page);
                 }
 
@@ -171,6 +187,14 @@ describe('Playwright Json', () => {
                 }
 
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
                   await jestPlaywright.saveCoverage(page);
                 }
 
@@ -251,6 +275,14 @@ describe('Playwright Json', () => {
                 }
 
                 if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
                   await jestPlaywright.saveCoverage(page);
                 }
 
@@ -323,172 +355,196 @@ describe('Playwright Json', () => {
         },
       };
       expect(transformPlaywrightJson(input)).toMatchInlineSnapshot(`
-              Object {
-                "example-header": "describe(\\"Example/Header\\", () => {
-                describe(\\"Logged In\\", () => {
-                  it(\\"test\\", async () => {
-                    const testFn = async () => {
-                      const context = {
-                        id: \\"example-header--logged-in\\",
-                        title: \\"Example/Header\\",
-                        name: \\"Logged In\\"
-                      };
-                      page.on('pageerror', err => {
-                        page.evaluate(({
-                          id,
-                          err
-                        }) => __throwError(id, err), {
-                          id: \\"example-header--logged-in\\",
-                          err: err.message
-                        });
-                      });
-
-                      if (global.__sbPreRender) {
-                        await global.__sbPreRender(page, context);
-                      }
-
-                      const result = await page.evaluate(({
-                        id,
-                        hasPlayFn
-                      }) => __test(id, hasPlayFn), {
-                        id: \\"example-header--logged-in\\"
-                      });
-
-                      if (global.__sbPostRender) {
-                        await global.__sbPostRender(page, context);
-                      }
-
-                      if (global.__sbCollectCoverage) {
-                        await jestPlaywright.saveCoverage(page);
-                      }
-
-                      return result;
-                    };
-
-                    try {
-                      await testFn();
-                    } catch (err) {
-                      if (err.toString().includes('Execution context was destroyed')) {
-                        await jestPlaywright.resetPage();
-                        await setupPage(global.page);
-                        await testFn();
-                      } else {
-                        throw err;
-                      }
-                    }
+        Object {
+          "example-header": "describe(\\"Example/Header\\", () => {
+          describe(\\"Logged In\\", () => {
+            it(\\"test\\", async () => {
+              const testFn = async () => {
+                const context = {
+                  id: \\"example-header--logged-in\\",
+                  title: \\"Example/Header\\",
+                  name: \\"Logged In\\"
+                };
+                page.on('pageerror', err => {
+                  page.evaluate(({
+                    id,
+                    err
+                  }) => __throwError(id, err), {
+                    id: \\"example-header--logged-in\\",
+                    err: err.message
                   });
                 });
-                describe(\\"Logged Out\\", () => {
-                  it(\\"test\\", async () => {
-                    const testFn = async () => {
-                      const context = {
-                        id: \\"example-header--logged-out\\",
-                        title: \\"Example/Header\\",
-                        name: \\"Logged Out\\"
-                      };
-                      page.on('pageerror', err => {
-                        page.evaluate(({
-                          id,
-                          err
-                        }) => __throwError(id, err), {
-                          id: \\"example-header--logged-out\\",
-                          err: err.message
-                        });
-                      });
-
-                      if (global.__sbPreRender) {
-                        await global.__sbPreRender(page, context);
-                      }
-
-                      const result = await page.evaluate(({
-                        id,
-                        hasPlayFn
-                      }) => __test(id, hasPlayFn), {
-                        id: \\"example-header--logged-out\\"
-                      });
-
-                      if (global.__sbPostRender) {
-                        await global.__sbPostRender(page, context);
-                      }
-
-                      if (global.__sbCollectCoverage) {
-                        await jestPlaywright.saveCoverage(page);
-                      }
-
-                      return result;
-                    };
-
-                    try {
-                      await testFn();
-                    } catch (err) {
-                      if (err.toString().includes('Execution context was destroyed')) {
-                        await jestPlaywright.resetPage();
-                        await setupPage(global.page);
-                        await testFn();
-                      } else {
-                        throw err;
-                      }
-                    }
+
+                if (global.__sbPreRender) {
+                  await global.__sbPreRender(page, context);
+                }
+
+                const result = await page.evaluate(({
+                  id,
+                  hasPlayFn
+                }) => __test(id, hasPlayFn), {
+                  id: \\"example-header--logged-in\\"
+                });
+
+                if (global.__sbPostRender) {
+                  await global.__sbPostRender(page, context);
+                }
+
+                if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
+                  await jestPlaywright.saveCoverage(page);
+                }
+
+                return result;
+              };
+
+              try {
+                await testFn();
+              } catch (err) {
+                if (err.toString().includes('Execution context was destroyed')) {
+                  await jestPlaywright.resetPage();
+                  await setupPage(global.page);
+                  await testFn();
+                } else {
+                  throw err;
+                }
+              }
+            });
+          });
+          describe(\\"Logged Out\\", () => {
+            it(\\"test\\", async () => {
+              const testFn = async () => {
+                const context = {
+                  id: \\"example-header--logged-out\\",
+                  title: \\"Example/Header\\",
+                  name: \\"Logged Out\\"
+                };
+                page.on('pageerror', err => {
+                  page.evaluate(({
+                    id,
+                    err
+                  }) => __throwError(id, err), {
+                    id: \\"example-header--logged-out\\",
+                    err: err.message
                   });
                 });
-              });",
-                "example-page": "describe(\\"Example/Page\\", () => {
-                describe(\\"Logged In\\", () => {
-                  it(\\"test\\", async () => {
-                    const testFn = async () => {
-                      const context = {
-                        id: \\"example-page--logged-in\\",
-                        title: \\"Example/Page\\",
-                        name: \\"Logged In\\"
-                      };
-                      page.on('pageerror', err => {
-                        page.evaluate(({
-                          id,
-                          err
-                        }) => __throwError(id, err), {
-                          id: \\"example-page--logged-in\\",
-                          err: err.message
-                        });
-                      });
-
-                      if (global.__sbPreRender) {
-                        await global.__sbPreRender(page, context);
-                      }
-
-                      const result = await page.evaluate(({
-                        id,
-                        hasPlayFn
-                      }) => __test(id, hasPlayFn), {
-                        id: \\"example-page--logged-in\\"
-                      });
-
-                      if (global.__sbPostRender) {
-                        await global.__sbPostRender(page, context);
-                      }
-
-                      if (global.__sbCollectCoverage) {
-                        await jestPlaywright.saveCoverage(page);
-                      }
-
-                      return result;
-                    };
-
-                    try {
-                      await testFn();
-                    } catch (err) {
-                      if (err.toString().includes('Execution context was destroyed')) {
-                        await jestPlaywright.resetPage();
-                        await setupPage(global.page);
-                        await testFn();
-                      } else {
-                        throw err;
-                      }
-                    }
+
+                if (global.__sbPreRender) {
+                  await global.__sbPreRender(page, context);
+                }
+
+                const result = await page.evaluate(({
+                  id,
+                  hasPlayFn
+                }) => __test(id, hasPlayFn), {
+                  id: \\"example-header--logged-out\\"
+                });
+
+                if (global.__sbPostRender) {
+                  await global.__sbPostRender(page, context);
+                }
+
+                if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
+                  await jestPlaywright.saveCoverage(page);
+                }
+
+                return result;
+              };
+
+              try {
+                await testFn();
+              } catch (err) {
+                if (err.toString().includes('Execution context was destroyed')) {
+                  await jestPlaywright.resetPage();
+                  await setupPage(global.page);
+                  await testFn();
+                } else {
+                  throw err;
+                }
+              }
+            });
+          });
+        });",
+          "example-page": "describe(\\"Example/Page\\", () => {
+          describe(\\"Logged In\\", () => {
+            it(\\"test\\", async () => {
+              const testFn = async () => {
+                const context = {
+                  id: \\"example-page--logged-in\\",
+                  title: \\"Example/Page\\",
+                  name: \\"Logged In\\"
+                };
+                page.on('pageerror', err => {
+                  page.evaluate(({
+                    id,
+                    err
+                  }) => __throwError(id, err), {
+                    id: \\"example-page--logged-in\\",
+                    err: err.message
                   });
                 });
-              });",
+
+                if (global.__sbPreRender) {
+                  await global.__sbPreRender(page, context);
+                }
+
+                const result = await page.evaluate(({
+                  id,
+                  hasPlayFn
+                }) => __test(id, hasPlayFn), {
+                  id: \\"example-page--logged-in\\"
+                });
+
+                if (global.__sbPostRender) {
+                  await global.__sbPostRender(page, context);
+                }
+
+                if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
+                  await jestPlaywright.saveCoverage(page);
+                }
+
+                return result;
+              };
+
+              try {
+                await testFn();
+              } catch (err) {
+                if (err.toString().includes('Execution context was destroyed')) {
+                  await jestPlaywright.resetPage();
+                  await setupPage(global.page);
+                  await testFn();
+                } else {
+                  throw err;
+                }
               }
-          `);
+            });
+          });
+        });",
+        }
+      `);
     });
 
     it('should skip docs-only stories', () => {
@@ -524,64 +580,72 @@ describe('Playwright Json', () => {
         },
       };
       expect(transformPlaywrightJson(input)).toMatchInlineSnapshot(`
-              Object {
-                "example-page": "describe(\\"Example/Page\\", () => {
-                describe(\\"Logged In\\", () => {
-                  it(\\"test\\", async () => {
-                    const testFn = async () => {
-                      const context = {
-                        id: \\"example-page--logged-in\\",
-                        title: \\"Example/Page\\",
-                        name: \\"Logged In\\"
-                      };
-                      page.on('pageerror', err => {
-                        page.evaluate(({
-                          id,
-                          err
-                        }) => __throwError(id, err), {
-                          id: \\"example-page--logged-in\\",
-                          err: err.message
-                        });
-                      });
-
-                      if (global.__sbPreRender) {
-                        await global.__sbPreRender(page, context);
-                      }
-
-                      const result = await page.evaluate(({
-                        id,
-                        hasPlayFn
-                      }) => __test(id, hasPlayFn), {
-                        id: \\"example-page--logged-in\\"
-                      });
-
-                      if (global.__sbPostRender) {
-                        await global.__sbPostRender(page, context);
-                      }
-
-                      if (global.__sbCollectCoverage) {
-                        await jestPlaywright.saveCoverage(page);
-                      }
-
-                      return result;
-                    };
-
-                    try {
-                      await testFn();
-                    } catch (err) {
-                      if (err.toString().includes('Execution context was destroyed')) {
-                        await jestPlaywright.resetPage();
-                        await setupPage(global.page);
-                        await testFn();
-                      } else {
-                        throw err;
-                      }
-                    }
+        Object {
+          "example-page": "describe(\\"Example/Page\\", () => {
+          describe(\\"Logged In\\", () => {
+            it(\\"test\\", async () => {
+              const testFn = async () => {
+                const context = {
+                  id: \\"example-page--logged-in\\",
+                  title: \\"Example/Page\\",
+                  name: \\"Logged In\\"
+                };
+                page.on('pageerror', err => {
+                  page.evaluate(({
+                    id,
+                    err
+                  }) => __throwError(id, err), {
+                    id: \\"example-page--logged-in\\",
+                    err: err.message
                   });
                 });
-              });",
+
+                if (global.__sbPreRender) {
+                  await global.__sbPreRender(page, context);
+                }
+
+                const result = await page.evaluate(({
+                  id,
+                  hasPlayFn
+                }) => __test(id, hasPlayFn), {
+                  id: \\"example-page--logged-in\\"
+                });
+
+                if (global.__sbPostRender) {
+                  await global.__sbPostRender(page, context);
+                }
+
+                if (global.__sbCollectCoverage) {
+                  const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
+
+                  if (!isCoverageSetupCorrectly) {
+                    throw new Error(\`[Test runner] An error occurred when evaluating code coverage:
+        The code in this story is not instrumented, which means the coverage setup is likely not correct.
+        More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage\`);
+                  }
+
+                  await jestPlaywright.saveCoverage(page);
+                }
+
+                return result;
+              };
+
+              try {
+                await testFn();
+              } catch (err) {
+                if (err.toString().includes('Execution context was destroyed')) {
+                  await jestPlaywright.resetPage();
+                  await setupPage(global.page);
+                  await testFn();
+                } else {
+                  throw err;
+                }
               }
-          `);
+            });
+          });
+        });",
+        }
+      `);
     });
   });
 });
diff --git a/src/setup-page.ts b/src/setup-page.ts
index 2a7f4b08..af907d57 100644
--- a/src/setup-page.ts
+++ b/src/setup-page.ts
@@ -48,19 +48,6 @@ export const setupPage = async (page: Page) => {
     throw err;
   });
 
-  if (isCoverageMode) {
-    const isCoverageSetupCorrectly = await page.evaluate(() => '__coverage__' in window);
-    if (!isCoverageSetupCorrectly) {
-      throw new Error(
-        dedent`
-          [Test runner] An error occurred when evaluating code coverage:
-          The code in Storybook is not instrumented, which means the coverage setup is not correct.
-          More info: https://github.com/storybookjs/test-runner#setting-up-code-coverage
-        `
-      );
-    }
-  }
-
   // if we ever want to log something from the browser to node
   await page.exposeBinding('logToPage', (_, message) => console.log(message));