feat(vitest): allow env to be stubbed to undefined (#6359)

Co-authored-by: Vladimir <sleuths.slews0s@icloud.com>
vitest-dev · Sep 3, 2024 · c3b2757 · c3b2757
1 parent ad6e72f
commit c3b2757
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/docs/api/vi.md b/docs/api/vi.md
@@ -453,7 +453,7 @@ And while it is possible to spy on exports in `jsdom` or other Node.js environme
 
 ### vi.stubEnv {#vi-stubenv}
 
-- **Type:** `(name: string, value: string) => Vitest`
+- **Type:** `<T extends string>(name: T, value: T extends "PROD" | "DEV" | "SSR" ? boolean : string | undefined) => Vitest`
 
 Changes the value of environmental variable on `process.env` and `import.meta.env`. You can restore its value by calling `vi.unstubAllEnvs`.
 
@@ -467,6 +467,12 @@ vi.stubEnv('NODE_ENV', 'production')
 
 process.env.NODE_ENV === 'production'
 import.meta.env.NODE_ENV === 'production'
+
+vi.stubEnv('NODE_ENV', undefined)
+
+process.env.NODE_ENV === undefined
+import.meta.env.NODE_ENV === undefined
+
 // doesn't change other envs
 import.meta.env.MODE === 'development'
 ```

diff --git a/packages/vitest/src/integrations/vi.ts b/packages/vitest/src/integrations/vi.ts
@@ -346,7 +346,7 @@ export interface VitestUtils {
    */
   stubEnv: <T extends string>(
     name: T,
-    value: T extends 'PROD' | 'DEV' | 'SSR' ? boolean : string
+    value: T extends 'PROD' | 'DEV' | 'SSR' ? boolean : string | undefined
   ) => VitestUtils
 
   /**
@@ -642,13 +642,16 @@ function createVitest(): VitestUtils {
       return utils
     },
 
-    stubEnv(name: string, value: string | boolean) {
+    stubEnv(name: string, value: string | boolean | undefined) {
       if (!_stubsEnv.has(name)) {
         _stubsEnv.set(name, process.env[name])
       }
       if (_envBooleans.includes(name)) {
         process.env[name] = value ? '1' : ''
       }
+      else if (value === undefined) {
+        delete process.env[name]
+      }
       else {
         process.env[name] = String(value)
       }

diff --git a/test/core/test/stubs.test.ts b/test/core/test/stubs.test.ts
@@ -126,4 +126,13 @@ describe('stubbing envs', () => {
     vi.stubEnv('MY_TEST_ENV', true)
     expect(import.meta.env.MY_TEST_ENV).toBe('true')
   })
+
+  it('stubs to undefined and restores env', () => {
+    vi.stubEnv('VITE_TEST_UPDATE_ENV', undefined)
+    expect(import.meta.env.VITE_TEST_UPDATE_ENV).toBeUndefined()
+    expect(process.env.VITE_TEST_UPDATE_ENV).toBeUndefined()
+    vi.unstubAllEnvs()
+    expect(import.meta.env.VITE_TEST_UPDATE_ENV).toBe('development')
+    expect(process.env.VITE_TEST_UPDATE_ENV).toBe('development')
+  })
 })