[aftereachrestoreenv] Add afterEachRestoreEnv (#870)

## Summary:
Often, we want to modify environment variables during testing. However, this is a change that can easily bleed across test cases and so we must also properly handle resetting the variables after each test case. This is a common pattern that we can abstract away into a utility function, which I have done here.

Now, tests can just call `afterEachRestoreEnv` and it will automatically restore the environment variables to their original values after each test case.

Tests can be specific about the environment variables they want to restore, or restore the entire environment.

Issue: XXX-XXXX

## Test plan:
`yarn test`
`yarn typecheck`

Author: somewhatabstract

Reviewers: somewhatabstract, jeresig

Required Reviewers:

Approved By: jeresig

Checks: ⌛ Test (macos-latest, 16.x), ✅ codecov/project, ✅ CodeQL, ✅ Lint, typecheck, and coverage check (ubuntu-latest, 16.x), ✅ gerald, ✅ Prime node_modules cache for primary configuration (ubuntu-latest, 16.x), ✅ Analyze (javascript), ⏭  dependabot

Pull Request URL: #870

.changeset/fluffy-geckos-repair.md

@@ -0,0 +1,5 @@
+---
+"@khanacademy/wonder-stuff-testing": major
+---
+
+New afterEachRestoreEnv function added to `jest` support in Wonder Stuff Testing for capture and restoration of node environment variables.

.vscode/settings.json

@@ -9,7 +9,6 @@
     ],
     "files.trimTrailingWhitespace": true,
     "typescript.tsdk": "node_modules/typescript/lib",
-    "typescript.tsserver.experimental.enableProjectDiagnostics": true,
     "typescript.format.enable": false,
     "typescript.validate.enable": true,
     "javascript.validate.enable": true,

packages/wonder-stuff-testing/package.json

@@ -22,6 +22,9 @@
     "dependencies": {
         "@khanacademy/wonder-stuff-core": "^1.5.1"
     },
+    "peerDependencies": {
+        "jest": "^29"
+    },
     "devDependencies": {
         "@khanacademy/ws-dev-build-settings": "^2.0.0"
     },

packages/wonder-stuff-testing/src/jest/__tests__/after-each-restore-env.test.ts

@@ -0,0 +1,146 @@
+import * as JestWrappers from "../internal/jest-wrappers";
+
+import {afterEachRestoreEnv} from "../after-each-restore-env";
+
+jest.mock("../internal/jest-wrappers");
+
+describe("#afterEachRestoreEnv", () => {
+    const EXISTS_1 = process.env.EXISTS_1;
+    const EXISTS_2 = process.env.EXISTS_2;
+
+    const ABSENT_1 = process.env.ABSENT_1;
+    const ABSENT_2 = process.env.ABSENT_2;
+
+    afterEach(() => {
+        // In case our tests misbehave, we still want to be good test denizens
+        // so we restore the environment variables we changed without using the
+        // code under test.
+        if (EXISTS_1 === undefined) {
+            delete process.env.EXISTS_1;
+        } else {
+            process.env.EXISTS_1 = EXISTS_1;
+        }
+
+        if (EXISTS_2 === undefined) {
+            delete process.env.EXISTS_2;
+        } else {
+            process.env.EXISTS_2 = EXISTS_2;
+        }
+
+        if (ABSENT_1 === undefined) {
+            delete process.env.ABSENT_1;
+        } else {
+            process.env.ABSENT_1 = ABSENT_1;
+        }
+
+        if (ABSENT_2 === undefined) {
+            delete process.env.ABSENT_2;
+        } else {
+            process.env.ABSENT_2 = ABSENT_2;
+        }
+    });
+
+    it("should register an afterEach callback", () => {
+        // Arrange
+        const afterEachSpy = jest
+            .spyOn(JestWrappers, "afterEach")
+            .mockImplementationOnce(() => {});
+
+        // Act
+        afterEachRestoreEnv("FOO");
+
+        // Assert
+        expect(afterEachSpy).toHaveBeenCalledWith(expect.any(Function));
+    });
+
+    describe("function passed to afterEach", () => {
+        it("should restore changed environment variables to their original values", () => {
+            // Arrange
+            const afterEachSpy = jest
+                .spyOn(JestWrappers, "afterEach")
+                .mockImplementationOnce(() => {});
+
+            // Make sure the env vars exist.
+            process.env.EXISTS_1 = "exists-1";
+            process.env.EXISTS_2 = "exists-2";
+
+            // Act
+            // Capture the state and set up the callback.
+            afterEachRestoreEnv();
+            const afterEachCallback: any = afterEachSpy.mock.calls.at(-1)?.[0];
+            // Change the state.
+            process.env.EXISTS_1 = "exists-1-changed";
+            process.env.EXISTS_2 = "exists-2-changed";
+            // Restore the state.
+            afterEachCallback();
+            const result = [process.env.EXISTS_1, process.env.EXISTS_2];
+
+            // Assert
+            expect(result).toEqual(["exists-1", "exists-2"]);
+        });
+
+        it("should delete environment variables that were not set before", () => {
+            // Arrange
+            const afterEachSpy = jest
+                .spyOn(JestWrappers, "afterEach")
+                .mockImplementationOnce(() => {});
+
+            // Make sure the env vars don't exist.
+            delete process.env.ABSENT_1;
+            delete process.env.ABSENT_2;
+
+            // Act
+            // Capture the state and set up the callback.
+            afterEachRestoreEnv();
+            const afterEachCallback: any = afterEachSpy.mock.calls.at(-1)?.[0];
+            // Change the state.
+            process.env.ABSENT_1 = "absent-1-set";
+            process.env.ABSENT_2 = "absent-2-set";
+            // Restore the state.
+            afterEachCallback();
+            const result = [process.env.ABSENT_1, process.env.ABSENT_2];
+
+            // Assert
+            expect(result).toEqual([undefined, undefined]);
+        });
+
+        it("should only restore the variables it was asked to restore", () => {
+            // Arrange
+            const afterEachSpy = jest
+                .spyOn(JestWrappers, "afterEach")
+                .mockImplementationOnce(() => {});
+
+            // Make sure the env vars don't exist or exist as we want.
+            process.env.EXISTS_1 = "exists-1";
+            process.env.EXISTS_2 = "exists-2";
+            delete process.env.ABSENT_1;
+            delete process.env.ABSENT_2;
+
+            // Act
+            // Capture the state and set up the callback.
+            afterEachRestoreEnv("ABSENT_1", "EXISTS_1");
+            const afterEachCallback: any = afterEachSpy.mock.calls.at(-1)?.[0];
+            // Change the state.
+            process.env.EXISTS_1 = "exists-1-changed";
+            process.env.EXISTS_2 = "exists-2-changed";
+            process.env.ABSENT_1 = "absent-1-set";
+            process.env.ABSENT_2 = "absent-2-set";
+            // Restore the state.
+            afterEachCallback();
+            const result = [
+                process.env.EXISTS_1,
+                process.env.EXISTS_2,
+                process.env.ABSENT_1,
+                process.env.ABSENT_2,
+            ];
+
+            // Assert
+            expect(result).toEqual([
+                "exists-1",
+                "exists-2-changed",
+                undefined,
+                "absent-2-set",
+            ]);
+        });
+    });
+});

packages/wonder-stuff-testing/src/jest/__tests__/wait.test.ts

@@ -2,7 +2,6 @@ import {wait, waitForAnimationFrame} from "../wait";
 import * as VerifyRealTimers from "../internal/verify-real-timers";
 import * as UnverifiedWait from "../internal/unverified-wait";
 
-jest.mock("../internal/assert-jest");
 jest.mock("../internal/verify-real-timers");
 jest.mock("../internal/unverified-wait");
 

packages/wonder-stuff-testing/src/jest/after-each-restore-env.ts

@@ -0,0 +1,56 @@
+import {afterEach} from "./internal/jest-wrappers";
+
+/**
+ * Restore the values of the given environment variables after each test.
+ *
+ * This captures the values of the given environment variables on invocation
+ * and then, after each test case, it restores those values.
+ *
+ * @param variableNames The names of the environment variables to restore.
+ */
+export const afterEachRestoreEnv = (
+    ...variableNames: ReadonlyArray<string>
+): void => {
+    // We capture the variables to restore. If none are given, we capture all.
+    const restoreAll = variableNames.length === 0;
+    const variablesToCapture = restoreAll
+        ? Object.keys(process.env)
+        : variableNames;
+    // We capture the current values on invocation.
+    const originalValues = variablesToCapture.reduce(
+        (acc: Record<string, string | undefined>, variableName: string) => {
+            acc[variableName] = process.env[variableName];
+            return acc;
+        },
+        {},
+    );
+
+    /**
+     * Restore the value of the given variable.
+     */
+    const restoreValue = (variableName: string, value: string | undefined) => {
+        if (value === undefined) {
+            delete process.env[variableName];
+        } else {
+            process.env[variableName] = value;
+        }
+    };
+
+    // Now, in the afterEach call, we restore the environment state.
+    afterEach(() => {
+        // If we are restoriing all variables then we cannot rely solely on
+        // what was captured, but instead we must also check what is currently
+        // in the environment that was not captured to make sure we delete it.
+        if (restoreAll) {
+            for (const variableName of Object.keys(process.env)) {
+                if (!(variableName in originalValues)) {
+                    delete process.env[variableName];
+                }
+            }
+        }
+
+        for (const [variableName, value] of Object.entries(originalValues)) {
+            restoreValue(variableName, value);
+        }
+    });
+};

packages/wonder-stuff-testing/src/jest/index.ts

@@ -1,2 +1,3 @@
+export {afterEachRestoreEnv} from "./after-each-restore-env";
 export {isolateModules} from "./isolate-modules";
 export {wait} from "./wait";

packages/wonder-stuff-testing/src/jest/internal/jest-wrappers.ts

@@ -0,0 +1,10 @@
+/* istanbul ignore file */
+import * as JestGlobals from "@jest/globals";
+
+/**
+ * Wrap the jest `afterEach` function.
+ *
+ * This makes it easy to mock this in our tests without jest getting upset.
+ */
+export const afterEach: typeof JestGlobals.afterEach = (...args) =>
+    JestGlobals.afterEach(...args);