feat(browser): allow injecting scripts (#5656)

sheremet-va · web-flow · commit 21e58bd890c4 · 2024-05-03T16:00:59.000+02:00
diff --git a/docs/config/index.md b/docs/config/index.md
@@ -1612,6 +1612,55 @@ This option has no effect on tests running inside Node.js.
 
 If you rely on spying on ES modules with `vi.spyOn`, you can enable this experimental feature to allow spying on module exports.
 
+#### browser.indexScripts <Version>1.6.0</Version> {#browser-indexscripts}
+
+- **Type:** `BrowserScript[]`
+- **Default:** `[]`
+
+Custom scripts that should be injected into the index HTML before test iframes are initiated. This HTML document only sets up iframes and doesn't actually import your code.
+
+The script `src` and `content` will be processed by Vite plugins. Script should be provided in the following shape:
+
+```ts
+export interface BrowserScript {
+  /**
+   * If "content" is provided and type is "module", this will be its identifier.
+   *
+   * If you are using TypeScript, you can add `.ts` extension here for example.
+   * @default `injected-${index}.js`
+   */
+  id?: string
+  /**
+   * JavaScript content to be injected. This string is processed by Vite plugins if type is "module".
+   *
+   * You can use `id` to give Vite a hint about the file extension.
+   */
+  content?: string
+  /**
+   * Path to the script. This value is resolved by Vite so it can be a node module or a file path.
+   */
+  src?: string
+  /**
+   * If the script should be loaded asynchronously.
+   */
+  async?: boolean
+  /**
+   * Script type.
+   * @default 'module'
+   */
+  type?: string
+}
+```
+
+#### browser.testerScripts <Version>1.6.0</Version> {#browser-testerscripts}
+
+- **Type:** `BrowserScript[]`
+- **Default:** `[]`
+
+Custom scripts that should be injected into the tester HTML before the tests environment is initiated. This is useful to inject polyfills required for Vitest browser implementation. It is recommended to use [`setupFiles`](#setupfiles) in almost all cases instead of this.
+
+The script `src` and `content` will be processed by Vite plugins.
+
 ### clearMocks
 
 - **Type:** `boolean`
diff --git a/packages/browser/src/client/index.html b/packages/browser/src/client/index.html
@@ -22,6 +22,7 @@
       }
     </style>
     <script>{__VITEST_INJECTOR__}</script>
+    {__VITEST_SCRIPTS__}
   </head>
   <body>
     <iframe id="vitest-ui" src=""></iframe>
diff --git a/packages/browser/src/client/tester.html b/packages/browser/src/client/tester.html
@@ -16,6 +16,7 @@
       }
     </style>
     <script>{__VITEST_INJECTOR__}</script>
+    {__VITEST_SCRIPTS__}
   </head>
   <body>
     <script type="module" src="/tester.ts"></script>
diff --git a/packages/browser/src/node/index.ts b/packages/browser/src/node/index.ts
@@ -1,17 +1,14 @@
 import { fileURLToPath } from 'node:url'
 import { readFile } from 'node:fs/promises'
-import { basename, resolve } from 'pathe'
+import { basename, join, resolve } from 'pathe'
 import sirv from 'sirv'
-import type { Plugin } from 'vite'
+import type { Plugin, ViteDevServer } from 'vite'
 import type { ResolvedConfig } from 'vitest'
-import type { WorkspaceProject } from 'vitest/node'
+import type { BrowserScript, WorkspaceProject } from 'vitest/node'
 import { coverageConfigDefaults } from 'vitest/config'
+import { slash } from '@vitest/utils'
 import { injectVitestModule } from './esmInjector'
 
-function replacer(code: string, values: Record<string, string>) {
-  return code.replace(/{\s*(\w+)\s*}/g, (_, key) => values[key] ?? '')
-}
-
 export default (project: WorkspaceProject, base = '/'): Plugin[] => {
   const pkgRoot = resolve(fileURLToPath(import.meta.url), '../..')
   const distRoot = resolve(pkgRoot, 'dist')
@@ -41,6 +38,8 @@ export default (project: WorkspaceProject, base = '/'): Plugin[] => {
           }
           next()
         })
+        let indexScripts: string | undefined
+        let testerScripts: string | undefined
         server.middlewares.use(async (req, res, next) => {
           if (!req.url)
             return next()
@@ -63,9 +62,13 @@ export default (project: WorkspaceProject, base = '/'): Plugin[] => {
           })
 
           if (url.pathname === base) {
+            if (!indexScripts)
+              indexScripts = await formatScripts(project.config.browser.indexScripts, server)
+
             const html = replacer(await runnerHtml, {
               __VITEST_FAVICON__: favicon,
               __VITEST_TITLE__: 'Vitest Browser Runner',
+              __VITEST_SCRIPTS__: indexScripts,
               __VITEST_INJECTOR__: injector,
             })
             res.write(html, 'utf-8')
@@ -77,9 +80,13 @@ export default (project: WorkspaceProject, base = '/'): Plugin[] => {
           // if decoded test file is "__vitest_all__" or not in the list of known files, run all tests
           const tests = decodedTestFile === '__vitest_all__' || !files.includes(decodedTestFile) ? '__vitest_browser_runner__.files' : JSON.stringify([decodedTestFile])
 
+          if (!testerScripts)
+            testerScripts = await formatScripts(project.config.browser.testerScripts, server)
+
           const html = replacer(await testerHtml, {
             __VITEST_FAVICON__: favicon,
             __VITEST_TITLE__: 'Vitest Browser Tester',
+            __VITEST_SCRIPTS__: testerScripts,
             __VITEST_INJECTOR__: injector,
             __VITEST_APPEND__:
             // TODO: have only a single global variable to not pollute the global scope
@@ -233,3 +240,22 @@ function wrapConfig(config: ResolvedConfig): ResolvedConfig {
         : undefined,
   }
 }
+
+function replacer(code: string, values: Record<string, string>) {
+  return code.replace(/{\s*(\w+)\s*}/g, (_, key) => values[key] ?? '')
+}
+
+async function formatScripts(scripts: BrowserScript[] | undefined, server: ViteDevServer) {
+  if (!scripts?.length)
+    return ''
+  const promises = scripts.map(async ({ content, src, async, id, type = 'module' }, index) => {
+    const srcLink = (src ? (await server.pluginContainer.resolveId(src))?.id : undefined) || src
+    const transformId = srcLink || join(server.config.root, `virtual__${id || `injected-${index}.js`}`)
+    await server.moduleGraph.ensureEntryFromUrl(transformId)
+    const contentProcessed = content && type === 'module'
+      ? (await server.pluginContainer.transform(content, transformId)).code
+      : content
+    return `<script type="${type}"${async ? ' async' : ''}${srcLink ? ` src="${slash(`/@fs/${srcLink}`)}"` : ''}>${contentProcessed || ''}</script>`
+  })
+  return (await Promise.all(promises)).join('\n')
+}
diff --git a/packages/vitest/src/node/cli/cli-config.ts b/packages/vitest/src/node/cli/cli-config.ts
@@ -350,6 +350,8 @@ export const cliOptionsConfig: VitestCLIOptions = {
       fileParallelism: {
         description: 'Should all test files run in parallel. Use `--browser.file-parallelism=false` to disable (default: same as `--file-parallelism`)',
       },
+      indexScripts: null,
+      testerScripts: null,
     },
   },
   pool: {
diff --git a/packages/vitest/src/node/index.ts b/packages/vitest/src/node/index.ts
@@ -13,4 +13,4 @@ export { VitestPackageInstaller } from './packageInstaller'
 export type { TestSequencer, TestSequencerConstructor } from './sequencers/types'
 export { BaseSequencer } from './sequencers/BaseSequencer'
 
-export type { BrowserProviderInitializationOptions, BrowserProvider, BrowserProviderOptions } from '../types/browser'
+export type { BrowserProviderInitializationOptions, BrowserProvider, BrowserProviderOptions, BrowserScript } from '../types/browser'
diff --git a/packages/vitest/src/types/browser.ts b/packages/vitest/src/types/browser.ts
@@ -93,6 +93,45 @@ export interface BrowserConfigOptions {
    * @default test.fileParallelism
    */
   fileParallelism?: boolean
+
+  /**
+   * Scripts injected into the tester iframe.
+   */
+  testerScripts?: BrowserScript[]
+
+  /**
+   * Scripts injected into the main window.
+   */
+  indexScripts?: BrowserScript[]
+}
+
+export interface BrowserScript {
+  /**
+   * If "content" is provided and type is "module", this will be its identifier.
+   *
+   * If you are using TypeScript, you can add `.ts` extension here for example.
+   * @default `injected-${index}.js`
+   */
+  id?: string
+  /**
+   * JavaScript content to be injected. This string is processed by Vite plugins if type is "module".
+   *
+   * You can use `id` to give Vite a hint about the file extension.
+   */
+  content?: string
+  /**
+   * Path to the script. This value is resolved by Vite so it can be a node module or a file path.
+   */
+  src?: string
+  /**
+   * If the script should be loaded asynchronously.
+   */
+  async?: boolean
+  /**
+   * Script type.
+   * @default 'module'
+   */
+  type?: string
 }
 
 export interface ResolvedBrowserOptions extends BrowserConfigOptions {
diff --git a/packages/vitest/src/types/config.ts b/packages/vitest/src/types/config.ts
@@ -16,6 +16,7 @@ import type { BenchmarkUserOptions } from './benchmark'
 import type { BrowserConfigOptions, ResolvedBrowserOptions } from './browser'
 import type { Pool, PoolOptions } from './pool-options'
 
+export type { BrowserScript, BrowserConfigOptions } from './browser'
 export type { SequenceHooks, SequenceSetupFiles } from '@vitest/runner'
 
 export type BuiltinEnvironment = 'node' | 'jsdom' | 'happy-dom' | 'edge-runtime'
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/test/browser/injected-lib/index.js b/test/browser/injected-lib/index.js
@@ -0,0 +1 @@
+__injected.push(4)
diff --git a/test/browser/injected-lib/package.json b/test/browser/injected-lib/package.json
@@ -0,0 +1,7 @@
+{
+  "name": "@vitest/injected-lib",
+  "type": "module",
+  "exports": {
+    "default": "./index.js"
+  }
+}
diff --git a/test/browser/injected.ts b/test/browser/injected.ts
@@ -0,0 +1,2 @@
+// @ts-expect-error not typed global
+;(__injected as string[]).push(3)
diff --git a/test/browser/package.json b/test/browser/package.json
@@ -15,6 +15,7 @@
     "@vitejs/plugin-basic-ssl": "^1.0.2",
     "@vitest/browser": "workspace:*",
     "@vitest/cjs-lib": "link:./cjs-lib",
+    "@vitest/injected-lib": "link:./injected-lib",
     "execa": "^7.1.1",
     "playwright": "^1.41.0",
     "url": "^0.11.3",
diff --git a/test/browser/specs/runner.test.ts b/test/browser/specs/runner.test.ts
@@ -1,4 +1,4 @@
-import { beforeAll, describe, expect, test } from 'vitest'
+import { beforeAll, describe, expect, onTestFailed, test } from 'vitest'
 import { runBrowserTests } from './utils'
 
 describe.each([
@@ -26,8 +26,12 @@ describe.each([
   })
 
   test(`[${description}] tests are actually running`, () => {
-    expect(browserResultJson.testResults).toHaveLength(14)
-    expect(passedTests).toHaveLength(12)
+    onTestFailed(() => {
+      console.error(stderr)
+    })
+
+    expect(browserResultJson.testResults).toHaveLength(15)
+    expect(passedTests).toHaveLength(13)
     expect(failedTests).toHaveLength(2)
 
     expect(stderr).not.toContain('has been externalized for browser compatibility')
diff --git a/test/browser/test/injected.test.ts b/test/browser/test/injected.test.ts
@@ -0,0 +1,10 @@
+import { expect, test } from 'vitest'
+
+test('injected values are correct', () => {
+  expect((globalThis as any).__injected).toEqual([
+    1,
+    2,
+    3,
+    4,
+  ])
+})
diff --git a/test/browser/vitest.config.mts b/test/browser/vitest.config.mts
@@ -29,6 +29,37 @@ export default defineConfig({
       provider,
       isolate: false,
       slowHijackESM: true,
+      testerScripts: [
+        {
+          content: 'globalThis.__injected = []',
+          type: 'text/javascript',
+        },
+        {
+          content: '__injected.push(1)',
+        },
+        {
+          id: 'ts.ts',
+          content: '(__injected as string[]).push(2)',
+        },
+        {
+          src: './injected.ts',
+        },
+        {
+          src: '@vitest/injected-lib',
+        },
+      ],
+      indexScripts: [
+        {
+          content: 'console.log("Hello, World");globalThis.__injected = []',
+          type: 'text/javascript',
+        },
+        {
+          content: 'import "./injected.ts"',
+        },
+        {
+          content: 'if(__injected[0] !== 3) throw new Error("injected not working")',
+        },
+      ],
     },
     alias: {
       '#src': resolve(dir, './src'),

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,7 @@`
`22`	`22`	`}`
`23`	`23`	`</style>`
`24`	`24`	`<script>{__VITEST_INJECTOR__}</script>`
	`25`	`+ {__VITEST_SCRIPTS__}`
`25`	`26`	`</head>`
`26`	`27`	`<body>`
`27`	`28`	`<iframe id="vitest-ui" src=""></iframe>`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@`
`16`	`16`	`}`
`17`	`17`	`</style>`
`18`	`18`	`<script>{__VITEST_INJECTOR__}</script>`
	`19`	`+ {__VITEST_SCRIPTS__}`
`19`	`20`	`</head>`
`20`	`21`	`<body>`
`21`	`22`	`<script type="module" src="/tester.ts"></script>`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+// @ts-expect-error not typed global`
	`2`	`+;(__injected as string[]).push(3)`