From 85d8c14b920bbe7291075bc03c09c915b8ec9cd3 Mon Sep 17 00:00:00 2001
From: Joe Previte <jjprevite@gmail.com>
Date: Mon, 26 Jul 2021 11:21:30 -0700
Subject: [PATCH] refactor: use Teffen's solution for useEnv

---
 test/unit/helpers.test.ts          | 29 +++++++++++++++++++-
 test/unit/node/proxy_agent.test.ts | 44 ++++++++----------------------
 test/utils/helpers.ts              | 24 ++++++++++++++++
 3 files changed, 64 insertions(+), 33 deletions(-)

diff --git a/test/unit/helpers.test.ts b/test/unit/helpers.test.ts
index 74485475..55187788 100644
--- a/test/unit/helpers.test.ts
+++ b/test/unit/helpers.test.ts
@@ -1,5 +1,5 @@
 import { promises as fs } from "fs"
-import { tmpdir } from "../../test/utils/helpers"
+import { tmpdir, useEnv } from "../../test/utils/helpers"
 
 /**
  * This file is for testing test helpers (not core code).
@@ -12,3 +12,30 @@ describe("test helpers", () => {
     expect(fs.access(pathToTempDir)).resolves.toStrictEqual(undefined)
   })
 })
+
+describe("useEnv", () => {
+  beforeAll(() => {
+    jest.resetModules()
+    process.env.TEST_USE_ENV = "test environment variable"
+  })
+  afterAll(() => {
+    delete process.env.TEST_USE_ENV
+  })
+  it("should set and reset the env var", () => {
+    const envKey = "TEST_ENV_VAR"
+    const [setValue, resetValue] = useEnv(envKey)
+    setValue("hello-world")
+    expect(process.env[envKey]).toEqual("hello-world")
+    resetValue()
+    expect(process.env[envKey]).toEqual(undefined)
+  })
+  it("should set and reset the env var where a value was already set", () => {
+    const envKey = "TEST_USE_ENV"
+    expect(process.env[envKey]).toEqual("test environment variable")
+    const [setValue, resetValue] = useEnv(envKey)
+    setValue("hello there")
+    expect(process.env[envKey]).toEqual("hello there")
+    resetValue()
+    expect(process.env[envKey]).toEqual("test environment variable")
+  })
+})
diff --git a/test/unit/node/proxy_agent.test.ts b/test/unit/node/proxy_agent.test.ts
index 9868c19f..a2552b7f 100644
--- a/test/unit/node/proxy_agent.test.ts
+++ b/test/unit/node/proxy_agent.test.ts
@@ -1,59 +1,39 @@
 import { shouldEnableProxy } from "../../../src/node/proxy_agent"
-
-/**
- * Helper function to set an env variable.
- *
- * Returns a function to cleanup the env variable.
- */
-function setEnv(name: string, value: string) {
-  process.env[name] = value
-
-  return {
-    cleanup() {
-      delete process.env[name]
-    },
-  }
-}
+import { useEnv } from "../../utils/helpers"
 
 describe("shouldEnableProxy", () => {
-  // Source: https://stackoverflow.com/a/48042799
-  const OLD_ENV = process.env
+  const [setHTTPProxy, resetHTTPProxy] = useEnv("HTTP_PROXY")
+  const [setHTTPSProxy, resetHTTPSProxy] = useEnv("HTTPS_PROXY")
+  const [setNoProxy, resetNoProxy] = useEnv("NO_PROXY")
 
   beforeEach(() => {
     jest.resetModules() // Most important - it clears the cache
-    process.env = { ...OLD_ENV } // Make a copy
-  })
-
-  afterAll(() => {
-    process.env = OLD_ENV // Restore old environment
+    resetHTTPProxy()
+    resetNoProxy()
+    resetHTTPSProxy()
   })
 
   it("returns true when HTTP_PROXY is set", () => {
-    const { cleanup } = setEnv("HTTP_PROXY", "http://proxy.example.com")
+    setHTTPProxy("http://proxy.example.com")
     expect(shouldEnableProxy()).toBe(true)
-    cleanup()
   })
   it("returns true when HTTPS_PROXY is set", () => {
-    const { cleanup } = setEnv("HTTPS_PROXY", "http://proxy.example.com")
+    setHTTPSProxy("https://proxy.example.com")
     expect(shouldEnableProxy()).toBe(true)
-    cleanup()
   })
   it("returns false when NO_PROXY is set", () => {
-    const { cleanup } = setEnv("NO_PROXY", "*")
+    setNoProxy("*")
     expect(shouldEnableProxy()).toBe(false)
-    cleanup()
   })
   it("should return false when neither HTTP_PROXY nor HTTPS_PROXY is set", () => {
     expect(shouldEnableProxy()).toBe(false)
   })
   it("should return false when NO_PROXY is set to https://example.com", () => {
-    const { cleanup } = setEnv("NO_PROXY", "https://example.com")
+    setNoProxy("https://example.com")
     expect(shouldEnableProxy()).toBe(false)
-    cleanup()
   })
   it("should return false when NO_PROXY is set to http://example.com", () => {
-    const { cleanup } = setEnv("NO_PROXY", "http://example.com")
+    setNoProxy("http://example.com")
     expect(shouldEnableProxy()).toBe(false)
-    cleanup()
   })
 })
diff --git a/test/utils/helpers.ts b/test/utils/helpers.ts
index a6a8d451..a7b7a6dd 100644
--- a/test/utils/helpers.ts
+++ b/test/utils/helpers.ts
@@ -37,3 +37,27 @@ export async function tmpdir(testName: string): Promise<string> {
 
   return await fs.mkdtemp(path.join(dir, `${testName}-`), { encoding: "utf8" })
 }
+
+/**
+ * @description Helper function to use an environment variable.
+ *
+ * @returns an array (similar to useState in React) with a function
+ * to set the value and reset the value
+ */
+export function useEnv(key: string): [(nextValue: string | undefined) => string | undefined, () => void] {
+  const initialValue = process.env[key]
+  const setValue = (nextValue: string | undefined) => (process.env[key] = nextValue)
+  // Node automatically converts undefined to string 'undefined'
+  // when assigning an environment variable.
+  // which is why we need to delete it if it's supposed to be undefined
+  // Source: https://stackoverflow.com/a/60147167
+  const resetValue = () => {
+    if (initialValue !== undefined) {
+      process.env[key] = initialValue
+    } else {
+      delete process.env[key]
+    }
+  }
+
+  return [setValue, resetValue]
+}