From c3b2757cd3d0c9c32a23be4ab49b4e48586c5e0f Mon Sep 17 00:00:00 2001
From: Juan Sanchez <11725911+JSanchezIO@users.noreply.github.com>
Date: Tue, 3 Sep 2024 02:22:20 -0700
Subject: [PATCH] feat(vitest): allow env to be stubbed to undefined (#6359)

Co-authored-by: Vladimir <sleuths.slews0s@icloud.com>
---
 docs/api/vi.md                         | 8 +++++++-
 packages/vitest/src/integrations/vi.ts | 7 +++++--
 test/core/test/stubs.test.ts           | 9 +++++++++
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/docs/api/vi.md b/docs/api/vi.md
index 8288d7683b74..5a967e62d11c 100644
--- 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
index c3f8250668c7..94719c039de2 100644
--- 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
index e2d8a99e21a3..9b8301ff3253 100644
--- 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')
+  })
 })