From 5f007889bbe7eb43ff945d772c43abcf84caf734 Mon Sep 17 00:00:00 2001
From: Vladimir <sheremet.va@icloud.com>
Date: Wed, 9 Feb 2022 19:52:55 +0300
Subject: [PATCH] feat: allow changing import.meta.env (#714)

---
 .../vitest/src/node/plugins/envReplacer.ts    | 30 +++++++++++++++++++
 packages/vitest/src/node/plugins/index.ts     | 19 ++++++++++++
 test/core/src/env.ts                          |  3 ++
 test/core/test/env.test.ts                    | 19 ++++++++++++
 test/core/types/env.d.ts                      | 12 ++++++++
 test/core/vitest.config.ts                    |  3 ++
 6 files changed, 86 insertions(+)
 create mode 100644 packages/vitest/src/node/plugins/envReplacer.ts
 create mode 100644 test/core/src/env.ts
 create mode 100644 test/core/test/env.test.ts
 create mode 100644 test/core/types/env.d.ts

diff --git a/packages/vitest/src/node/plugins/envReplacer.ts b/packages/vitest/src/node/plugins/envReplacer.ts
new file mode 100644
index 000000000000..28c7027e182d
--- /dev/null
+++ b/packages/vitest/src/node/plugins/envReplacer.ts
@@ -0,0 +1,30 @@
+import MagicString from 'magic-string'
+import type { Plugin } from 'vite'
+
+export const EnvReplacerPlugin = (): Plugin => {
+  return {
+    name: 'vitest:env-replacer',
+    enforce: 'pre',
+    transform(code) {
+      let s: MagicString | null = null
+
+      const envs = code.matchAll(/\bimport\.meta\.env\b/g)
+
+      for (const env of envs) {
+        s ||= new MagicString(code)
+
+        const startIndex = env.index!
+        const endIndex = startIndex + env[0].length
+
+        s.overwrite(startIndex, endIndex, 'process.env')
+      }
+
+      if (s) {
+        return {
+          code: s.toString(),
+          map: s.generateMap({ hires: true }),
+        }
+      }
+    },
+  }
+}
diff --git a/packages/vitest/src/node/plugins/index.ts b/packages/vitest/src/node/plugins/index.ts
index 262cf4d52d8d..b5f1f0b24dc9 100644
--- a/packages/vitest/src/node/plugins/index.ts
+++ b/packages/vitest/src/node/plugins/index.ts
@@ -6,6 +6,7 @@ import { resolveApiConfig } from '../config'
 import { Vitest } from '../core'
 import { GlobalSetupPlugin } from './globalSetup'
 import { MocksPlugin } from './mock'
+import { EnvReplacerPlugin } from './envReplacer'
 
 export async function VitestPlugin(options: UserConfig = {}, ctx = new Vitest()): Promise<VitePlugin[]> {
   let haveStarted = false
@@ -57,6 +58,23 @@ export async function VitestPlugin(options: UserConfig = {}, ctx = new Vitest())
         )
         options.api = resolveApiConfig(options)
         options.watch = options.watch && !options.run
+
+        process.env.BASE_URL ??= viteConfig.base
+        process.env.MODE ??= viteConfig.mode
+        // process.env can have only string values and will cast string on it if we pass other type,
+        // so we are making them truthy
+        process.env.PROD ??= viteConfig.env.PROD ? '1' : ''
+        process.env.DEV ??= viteConfig.env.DEV ? '1' : ''
+        process.env.SSR ??= '1'
+
+        // account for user env defines
+        for (const key in viteConfig.define) {
+          if (key.startsWith('import.meta.env.')) {
+            const val = viteConfig.define[key]
+            const envKey = key.slice('import.meta.env.'.length)
+            process.env[envKey] = typeof val === 'string' ? JSON.parse(val) : val
+          }
+        }
       },
       async configureServer(server) {
         if (haveStarted)
@@ -71,6 +89,7 @@ export async function VitestPlugin(options: UserConfig = {}, ctx = new Vitest())
           await server.watcher.close()
       },
     },
+    EnvReplacerPlugin(),
     MocksPlugin(),
     GlobalSetupPlugin(ctx),
     options.ui
diff --git a/test/core/src/env.ts b/test/core/src/env.ts
new file mode 100644
index 000000000000..dc81c1f038c2
--- /dev/null
+++ b/test/core/src/env.ts
@@ -0,0 +1,3 @@
+export function getAuthToken() {
+  return import.meta.env.AUTH_TOKEN
+}
diff --git a/test/core/test/env.test.ts b/test/core/test/env.test.ts
new file mode 100644
index 000000000000..fee811c1c156
--- /dev/null
+++ b/test/core/test/env.test.ts
@@ -0,0 +1,19 @@
+import { expect, test } from 'vitest'
+import { getAuthToken } from '../src/env'
+
+test('can reassign env locally', () => {
+  import.meta.env.VITEST_ENV = 'TEST'
+  expect(import.meta.env.VITEST_ENV).toBe('TEST')
+})
+
+test('can reassign env everywhere', () => {
+  import.meta.env.AUTH_TOKEN = '123'
+  expect(getAuthToken()).toBe('123')
+  process.env.AUTH_TOKEN = '321'
+  expect(getAuthToken()).toBe('321')
+})
+
+test('can see env in "define"', () => {
+  expect(import.meta.env.TEST_NAME).toBe('hello world')
+  expect(process.env.TEST_NAME).toBe('hello world')
+})
diff --git a/test/core/types/env.d.ts b/test/core/types/env.d.ts
new file mode 100644
index 000000000000..c45958f8bec8
--- /dev/null
+++ b/test/core/types/env.d.ts
@@ -0,0 +1,12 @@
+interface ImportMeta {
+  readonly env: ImportMetaEnv
+}
+
+interface ImportMetaEnv {
+  [key: string]: string | boolean | undefined
+  BASE_URL: string
+  MODE: string
+  DEV: boolean
+  PROD: boolean
+  SSR: boolean
+}
diff --git a/test/core/vitest.config.ts b/test/core/vitest.config.ts
index c53eaf09ce5c..3a2f0b1749d7 100644
--- a/test/core/vitest.config.ts
+++ b/test/core/vitest.config.ts
@@ -17,6 +17,9 @@ export default defineConfig({
       },
     },
   ],
+  define: {
+    'import.meta.env.TEST_NAME': '"hello world"',
+  },
   test: {
     testTimeout: 2000,
     // threads: false,