feat(vitest): allow overiding package installer with public API (#4936)

sheremet-va · web-flow · commit c2cceebbfd28 · 2024-01-12T17:04:31.000+01:00
diff --git a/packages/browser/src/node/providers/playwright.ts b/packages/browser/src/node/providers/playwright.ts
@@ -27,8 +27,8 @@ export class PlaywrightBrowserProvider implements BrowserProvider {
     return playwrightBrowsers
   }
 
-  async initialize(ctx: WorkspaceProject, { browser, options }: PlaywrightProviderOptions) {
-    this.ctx = ctx
+  initialize(project: WorkspaceProject, { browser, options }: PlaywrightProviderOptions) {
+    this.ctx = project
     this.browser = browser
     this.options = options as any
   }
diff --git a/packages/vitest/src/integrations/browser.ts b/packages/vitest/src/integrations/browser.ts
@@ -1,17 +1,12 @@
-import { ensurePackageInstalled } from '../node/pkg'
+import type { WorkspaceProject } from '../node/workspace'
 import type { BrowserProviderModule, ResolvedBrowserOptions } from '../types/browser'
 
-interface Loader {
-  root: string
-  executeId: (id: string) => any
-}
-
 const builtinProviders = ['webdriverio', 'playwright', 'none']
 
-export async function getBrowserProvider(options: ResolvedBrowserOptions, loader: Loader): Promise<BrowserProviderModule> {
+export async function getBrowserProvider(options: ResolvedBrowserOptions, project: WorkspaceProject): Promise<BrowserProviderModule> {
   if (options.provider == null || builtinProviders.includes(options.provider)) {
-    await ensurePackageInstalled('@vitest/browser', loader.root)
-    const providers = await loader.executeId('@vitest/browser/providers') as {
+    await project.ctx.packageInstaller.ensureInstalled('@vitest/browser', project.config.root)
+    const providers = await project.runner.executeId('@vitest/browser/providers') as {
       webdriverio: BrowserProviderModule
       playwright: BrowserProviderModule
       none: BrowserProviderModule
@@ -23,7 +18,7 @@ export async function getBrowserProvider(options: ResolvedBrowserOptions, loader
   let customProviderModule
 
   try {
-    customProviderModule = await loader.executeId(options.provider) as { default: BrowserProviderModule }
+    customProviderModule = await project.runner.executeId(options.provider) as { default: BrowserProviderModule }
   }
   catch (error) {
     throw new Error(`Failed to load custom BrowserProvider from ${options.provider}`, { cause: error })
diff --git a/packages/vitest/src/integrations/browser/server.ts b/packages/vitest/src/integrations/browser/server.ts
@@ -1,6 +1,5 @@
 import { createServer } from 'vite'
 import { defaultBrowserPort } from '../../constants'
-import { ensurePackageInstalled } from '../../node/pkg'
 import { resolveApiServerConfig } from '../../node/config'
 import { CoverageTransform } from '../../node/plugins/coverageTransform'
 import type { WorkspaceProject } from '../../node/workspace'
@@ -10,7 +9,7 @@ import { resolveFsAllow } from '../../node/plugins/utils'
 export async function createBrowserServer(project: WorkspaceProject, configFile: string | undefined) {
   const root = project.config.root
 
-  await ensurePackageInstalled('@vitest/browser', root)
+  await project.ctx.packageInstaller.ensureInstalled('@vitest/browser', root)
 
   const configPath = typeof configFile === 'string' ? configFile : false
 
diff --git a/packages/vitest/src/node/cli-api.ts b/packages/vitest/src/node/cli-api.ts
@@ -4,9 +4,9 @@ import { EXIT_CODE_RESTART } from '../constants'
 import { CoverageProviderMap } from '../integrations/coverage'
 import { getEnvPackageName } from '../integrations/env'
 import type { UserConfig, Vitest, VitestRunMode } from '../types'
-import { ensurePackageInstalled } from './pkg'
 import { createVitest } from './create'
 import { registerConsoleShortcuts } from './stdin'
+import type { VitestOptions } from './core'
 
 export interface CliOptions extends UserConfig {
   /**
@@ -25,6 +25,7 @@ export async function startVitest(
   cliFilters: string[] = [],
   options: CliOptions = {},
   viteOverrides?: ViteUserConfig,
+  vitestOptions?: VitestOptions,
 ): Promise<Vitest | undefined> {
   process.env.TEST = 'true'
   process.env.VITEST = 'true'
@@ -60,14 +61,14 @@ export async function startVitest(
     options.typecheck.enabled = true
   }
 
-  const ctx = await createVitest(mode, options, viteOverrides)
+  const ctx = await createVitest(mode, options, viteOverrides, vitestOptions)
 
   if (mode === 'test' && ctx.config.coverage.enabled) {
     const provider = ctx.config.coverage.provider || 'v8'
     const requiredPackages = CoverageProviderMap[provider]
 
     if (requiredPackages) {
-      if (!await ensurePackageInstalled(requiredPackages, root)) {
+      if (!await ctx.packageInstaller.ensureInstalled(requiredPackages, root)) {
         process.exitCode = 1
         return ctx
       }
@@ -76,7 +77,7 @@ export async function startVitest(
 
   const environmentPackage = getEnvPackageName(ctx.config.environment)
 
-  if (environmentPackage && !await ensurePackageInstalled(environmentPackage, root)) {
+  if (environmentPackage && !await ctx.packageInstaller.ensureInstalled(environmentPackage, root)) {
     process.exitCode = 1
     return ctx
   }
diff --git a/packages/vitest/src/node/core.ts b/packages/vitest/src/node/core.ts
@@ -23,9 +23,14 @@ import { resolveConfig } from './config'
 import { Logger } from './logger'
 import { VitestCache } from './cache'
 import { WorkspaceProject, initializeProject } from './workspace'
+import { VitestPackageInstaller } from './packageInstaller'
 
 const WATCHER_DEBOUNCE = 100
 
+export interface VitestOptions {
+  packageInstaller?: VitestPackageInstaller
+}
+
 export class Vitest {
   config: ResolvedConfig = undefined!
   configOverride: Partial<ResolvedConfig> = {}
@@ -53,6 +58,8 @@ export class Vitest {
   restartsCount = 0
   runner: ViteNodeRunner = undefined!
 
+  public packageInstaller: VitestPackageInstaller
+
   private coreWorkspaceProject!: WorkspaceProject
 
   private resolvedProjects: WorkspaceProject[] = []
@@ -63,8 +70,10 @@ export class Vitest {
 
   constructor(
     public readonly mode: VitestRunMode,
+    options: VitestOptions = {},
   ) {
     this.logger = new Logger(this)
+    this.packageInstaller = options.packageInstaller || new VitestPackageInstaller()
   }
 
   private _onRestartListeners: OnServerRestartHandler[] = []
@@ -139,7 +148,7 @@ export class Vitest {
 
     this.reporters = resolved.mode === 'benchmark'
       ? await createBenchmarkReporters(toArray(resolved.benchmark?.reporters), this.runner)
-      : await createReporters(resolved.reporters, this.runner)
+      : await createReporters(resolved.reporters, this)
 
     this.cache.results.setConfig(resolved.root, resolved.cache)
     try {
diff --git a/packages/vitest/src/node/create.ts b/packages/vitest/src/node/create.ts
@@ -4,12 +4,13 @@ import type { InlineConfig as ViteInlineConfig, UserConfig as ViteUserConfig } f
 import { findUp } from 'find-up'
 import type { UserConfig, VitestRunMode } from '../types'
 import { configFiles } from '../constants'
+import type { VitestOptions } from './core'
 import { Vitest } from './core'
 import { VitestPlugin } from './plugins'
 import { createViteServer } from './vite'
 
-export async function createVitest(mode: VitestRunMode, options: UserConfig, viteOverrides: ViteUserConfig = {}) {
-  const ctx = new Vitest(mode)
+export async function createVitest(mode: VitestRunMode, options: UserConfig, viteOverrides: ViteUserConfig = {}, vitestOptions: VitestOptions = {}) {
+  const ctx = new Vitest(mode, vitestOptions)
   const root = resolve(options.root || process.cwd())
 
   const configPath = options.config === false
diff --git a/packages/vitest/src/node/index.ts b/packages/vitest/src/node/index.ts
@@ -7,6 +7,7 @@ export { registerConsoleShortcuts } from './stdin'
 export type { GlobalSetupContext } from './globalSetup'
 export type { WorkspaceSpec, ProcessPool } from './pool'
 export { createMethodsRPC } from './pools/rpc'
+export { VitestPackageInstaller } from './packageInstaller'
 
 export type { TestSequencer, TestSequencerConstructor } from './sequencers/types'
 export { BaseSequencer } from './sequencers/BaseSequencer'
diff --git a/packages/vitest/src/node/packageInstaller.ts b/packages/vitest/src/node/packageInstaller.ts
@@ -0,0 +1,52 @@
+import url from 'node:url'
+import { createRequire } from 'node:module'
+import c from 'picocolors'
+import { isPackageExists } from 'local-pkg'
+import { EXIT_CODE_RESTART } from '../constants'
+import { isCI } from '../utils/env'
+
+const __dirname = url.fileURLToPath(new URL('.', import.meta.url))
+
+export class VitestPackageInstaller {
+  async ensureInstalled(dependency: string, root: string) {
+    if (process.env.VITEST_SKIP_INSTALL_CHECKS)
+      return true
+
+    if (process.versions.pnp) {
+      const targetRequire = createRequire(__dirname)
+      try {
+        targetRequire.resolve(dependency, { paths: [root, __dirname] })
+        return true
+      }
+      catch (error) {
+      }
+    }
+
+    if (isPackageExists(dependency, { paths: [root, __dirname] }))
+      return true
+
+    const promptInstall = !isCI && process.stdout.isTTY
+
+    process.stderr.write(c.red(`${c.inverse(c.red(' MISSING DEPENDENCY '))} Cannot find dependency '${dependency}'\n\n`))
+
+    if (!promptInstall)
+      return false
+
+    const prompts = await import('prompts')
+    const { install } = await prompts.prompt({
+      type: 'confirm',
+      name: 'install',
+      message: c.reset(`Do you want to install ${c.green(dependency)}?`),
+    })
+
+    if (install) {
+      await (await import('@antfu/install-pkg')).installPackage(dependency, { dev: true })
+      // TODO: somehow it fails to load the package after installation, remove this when it's fixed
+      process.stderr.write(c.yellow(`\nPackage ${dependency} installed, re-run the command to start.\n`))
+      process.exit(EXIT_CODE_RESTART)
+      return true
+    }
+
+    return false
+  }
+}
diff --git a/packages/vitest/src/node/pkg.ts b/packages/vitest/src/node/pkg.ts
diff --git a/packages/vitest/src/node/plugins/index.ts b/packages/vitest/src/node/plugins/index.ts
@@ -3,7 +3,6 @@ import { relative } from 'pathe'
 import { configDefaults } from '../../defaults'
 import type { ResolvedConfig, UserConfig } from '../../types'
 import { deepMerge, notNullish, removeUndefinedValues, toArray } from '../../utils'
-import { ensurePackageInstalled } from '../pkg'
 import { resolveApiServerConfig } from '../config'
 import { Vitest } from '../core'
 import { generateScopedClassName } from '../../integrations/css/css-modules'
@@ -22,7 +21,7 @@ export async function VitestPlugin(options: UserConfig = {}, ctx = new Vitest('t
   const getRoot = () => ctx.config?.root || options.root || process.cwd()
 
   async function UIPlugin() {
-    await ensurePackageInstalled('@vitest/ui', getRoot())
+    await ctx.packageInstaller.ensureInstalled('@vitest/ui', getRoot())
     return (await import('@vitest/ui')).default(ctx)
   }
 
diff --git a/packages/vitest/src/node/reporters/utils.ts b/packages/vitest/src/node/reporters/utils.ts
@@ -1,6 +1,5 @@
 import type { ViteNodeRunner } from 'vite-node/client'
-import type { Reporter } from '../../types'
-import { ensurePackageInstalled } from '../pkg'
+import type { Reporter, Vitest } from '../../types'
 import { BenchmarkReportsMap, ReportersMap } from './index'
 import type { BenchmarkBuiltinReporters, BuiltinReporters } from './index'
 
@@ -19,11 +18,12 @@ async function loadCustomReporterModule<C extends Reporter>(path: string, runner
   return customReporterModule.default
 }
 
-function createReporters(reporterReferences: Array<string | Reporter | BuiltinReporters>, runner: ViteNodeRunner) {
+function createReporters(reporterReferences: Array<string | Reporter | BuiltinReporters>, ctx: Vitest) {
+  const runner = ctx.runner
   const promisedReporters = reporterReferences.map(async (referenceOrInstance) => {
     if (typeof referenceOrInstance === 'string') {
       if (referenceOrInstance === 'html') {
-        await ensurePackageInstalled('@vitest/ui', runner.root)
+        await ctx.packageInstaller.ensureInstalled('@vitest/ui', runner.root)
         const CustomReporter = await loadCustomReporterModule('@vitest/ui/reporter', runner)
         return new CustomReporter()
       }
diff --git a/packages/vitest/src/node/workspace.ts b/packages/vitest/src/node/workspace.ts
@@ -386,7 +386,7 @@ export class WorkspaceProject {
       return
     if (this.browserProvider)
       return
-    const Provider = await getBrowserProvider(this.config.browser, this.runner)
+    const Provider = await getBrowserProvider(this.config.browser, this)
     this.browserProvider = new Provider()
     const browser = this.config.browser.name
     const supportedBrowsers = this.browserProvider.getSupportedBrowsers()
diff --git a/packages/vitest/src/typecheck/typechecker.ts b/packages/vitest/src/typecheck/typechecker.ts
@@ -6,8 +6,7 @@ import { basename, extname, resolve } from 'pathe'
 import { TraceMap, generatedPositionFor } from '@vitest/utils/source-map'
 import type { RawSourceMap } from '@ampproject/remapping'
 import { getTasks } from '../utils'
-import { ensurePackageInstalled } from '../node/pkg'
-import type { Awaitable, File, ParsedStack, Task, TaskResultPack, TaskState, TscErrorInfo } from '../types'
+import type { Awaitable, File, ParsedStack, Task, TaskResultPack, TaskState, TscErrorInfo, Vitest } from '../types'
 import type { WorkspaceProject } from '../node/workspace'
 import { getRawErrsMapFromTsCompile, getTsconfig } from './parse'
 import { createIndexMap } from './utils'
@@ -225,16 +224,15 @@ export class Typechecker {
     this.process?.kill()
   }
 
-  protected async ensurePackageInstalled(root: string, checker: string) {
+  protected async ensurePackageInstalled(ctx: Vitest, checker: string) {
     if (checker !== 'tsc' && checker !== 'vue-tsc')
       return
     const packageName = checker === 'tsc' ? 'typescript' : 'vue-tsc'
-    await ensurePackageInstalled(packageName, root)
+    await ctx.packageInstaller.ensureInstalled(packageName, ctx.config.root)
   }
 
   public async prepare() {
     const { root, typecheck } = this.ctx.config
-    await this.ensurePackageInstalled(root, typecheck.checker)
 
     const { config, path } = await getTsconfig(root, typecheck)
 
diff --git a/test/reporters/tests/utils.test.ts b/test/reporters/tests/utils.test.ts
@@ -3,6 +3,7 @@
  */
 import { resolve } from 'pathe'
 import type { ViteNodeRunner } from 'vite-node/client'
+import type { Vitest } from 'vitest'
 import { describe, expect, test } from 'vitest'
 import { createReporters } from '../../../packages/vitest/src/node/reporters/utils'
 import { DefaultReporter } from '../../../packages/vitest/src/node/reporters/default'
@@ -12,29 +13,32 @@ const customReporterPath = resolve(__dirname, '../src/custom-reporter.js')
 const fetchModule = {
   executeId: (id: string) => import(id),
 } as ViteNodeRunner
+const ctx = {
+  runner: fetchModule,
+} as Vitest
 
 describe('Reporter Utils', () => {
   test('passing an empty array returns nothing', async () => {
-    const promisedReporters = await createReporters([], fetchModule)
+    const promisedReporters = await createReporters([], ctx)
     expect(promisedReporters).toHaveLength(0)
   })
 
   test('passing the name of a single built-in reporter returns a new instance', async () => {
-    const promisedReporters = await createReporters(['default'], fetchModule)
+    const promisedReporters = await createReporters(['default'], ctx)
     expect(promisedReporters).toHaveLength(1)
     const reporter = promisedReporters[0]
     expect(reporter).toBeInstanceOf(DefaultReporter)
   })
 
   test('passing in the path to a custom reporter returns a new instance', async () => {
-    const promisedReporters = await createReporters(([customReporterPath]), fetchModule)
+    const promisedReporters = await createReporters(([customReporterPath]), ctx)
     expect(promisedReporters).toHaveLength(1)
     const customReporter = promisedReporters[0]
     expect(customReporter).toBeInstanceOf(TestReporter)
   })
 
   test('passing in a mix of built-in and custom reporters works', async () => {
-    const promisedReporters = await createReporters(['default', customReporterPath], fetchModule)
+    const promisedReporters = await createReporters(['default', customReporterPath], ctx)
     expect(promisedReporters).toHaveLength(2)
     const defaultReporter = promisedReporters[0]
     expect(defaultReporter).toBeInstanceOf(DefaultReporter)

Original file line number	Diff line number	Diff line change
`@@ -27,8 +27,8 @@ export class PlaywrightBrowserProvider implements BrowserProvider {`
`27`	`27`	`return playwrightBrowsers`
`28`	`28`	`}`
`29`	`29`
`30`		`- async initialize(ctx: WorkspaceProject, { browser, options }: PlaywrightProviderOptions) {`
`31`		`- this.ctx = ctx`
	`30`	`+ initialize(project: WorkspaceProject, { browser, options }: PlaywrightProviderOptions) {`
	`31`	`+ this.ctx = project`
`32`	`32`	`this.browser = browser`
`33`	`33`	`this.options = options as any`
`34`	`34`	`}`