From 1dc08f58e8fe964fe7dd13c01fdea2fd0ef8d406 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cristian=20Pallar=C3=A9s?= <cristianp@wing.cloud>
Date: Thu, 15 Aug 2024 15:32:46 +0200
Subject: [PATCH] fix(cli): load env files depepnding on the wing command
 (#7021)

---
 .../console/server/src/utils/compiler.ts      |  2 +-
 apps/wing/src/commands/compile.ts             |  5 +++-
 apps/wing/src/commands/pack.ts                |  2 +-
 apps/wing/src/commands/run.ts                 |  2 +-
 libs/wingsdk/src/helpers.ts                   | 22 +++++++++++++---
 libs/wingsdk/test/helpers.test.ts             | 26 +++++++++++++++++++
 6 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/apps/wing-console/console/server/src/utils/compiler.ts b/apps/wing-console/console/server/src/utils/compiler.ts
index 442b8500334..c0f7ec0a8dc 100644
--- a/apps/wing-console/console/server/src/utils/compiler.ts
+++ b/apps/wing-console/console/server/src/utils/compiler.ts
@@ -53,7 +53,7 @@ export const createCompiler = ({
       return;
     }
 
-    loadEnvVariables({ cwd: dirname });
+    loadEnvVariables({ modes: ["run", "it"], cwd: dirname });
 
     isCompiling = true;
     await events.emit("compiling");
diff --git a/apps/wing/src/commands/compile.ts b/apps/wing/src/commands/compile.ts
index 4f2416382ea..67277344e3c 100644
--- a/apps/wing/src/commands/compile.ts
+++ b/apps/wing/src/commands/compile.ts
@@ -84,7 +84,10 @@ export async function compile(entrypoint?: string, options?: CompileOptions): Pr
     }
     entrypoint = wingFiles[0];
   }
-  loadEnvVariables({ cwd: resolve(dirname(entrypoint)) });
+  loadEnvVariables({
+    modes: options?.testing ? ["test"] : ["compile"],
+    cwd: resolve(dirname(entrypoint)),
+  });
   const coloring = chalk.supportsColor ? chalk.supportsColor.hasBasic : false;
   const compileOutput = await wingCompiler.compile(entrypoint, {
     ...options,
diff --git a/apps/wing/src/commands/pack.ts b/apps/wing/src/commands/pack.ts
index 21df640fb8b..065f2249ad6 100644
--- a/apps/wing/src/commands/pack.ts
+++ b/apps/wing/src/commands/pack.ts
@@ -80,7 +80,7 @@ export async function pack(options: PackageOptions = {}): Promise<string> {
   const outfile = options.outFile ? resolve(options.outFile) : undefined;
   const outdir = outfile ? path.dirname(outfile) : userDir;
 
-  loadEnvVariables({ cwd: userDir });
+  loadEnvVariables({ modes: ["pack"], cwd: userDir });
   // check package.json exists
   const originalPkgJsonPath = path.join(userDir, "package.json");
   if (!(await exists(originalPkgJsonPath))) {
diff --git a/apps/wing/src/commands/run.ts b/apps/wing/src/commands/run.ts
index 770109a11ef..122c676c577 100644
--- a/apps/wing/src/commands/run.ts
+++ b/apps/wing/src/commands/run.ts
@@ -73,7 +73,7 @@ export async function run(entrypoint?: string, options?: RunOptions) {
     throw new Error(entrypoint + " doesn't exist");
   }
 
-  loadEnvVariables({ cwd: resolve(dirname(entrypoint)) });
+  loadEnvVariables({ modes: ["run", "it"], cwd: resolve(dirname(entrypoint)) });
 
   if (options?.platform && options?.platform[0] !== BuiltinPlatform.SIM) {
     throw new Error(
diff --git a/libs/wingsdk/src/helpers.ts b/libs/wingsdk/src/helpers.ts
index 4e3a971a710..58cb31b4662 100644
--- a/libs/wingsdk/src/helpers.ts
+++ b/libs/wingsdk/src/helpers.ts
@@ -233,8 +233,6 @@ export function preflightClassSingleton(
   return root.resourceSingletons[type];
 }
 
-const DEFAULT_ENV_FILES = [`.env`, `.env.local`];
-
 /**
  * Options for loading environment variables.
  */
@@ -244,16 +242,32 @@ export interface EnvLoadOptions {
    * @default process.cwd()
    */
   readonly cwd?: string;
+
+  /**
+   * The modes to load the environment variables for.
+   */
+  readonly modes: string[];
 }
 
 /**
- * Loads environment variables from `.env` and `.env.local` files.
+ * Loads environment variables from:
+ * - `.env`
+ * - `.env.local`
+ * - `.env.{mode}`
+ * - `.env.{mode}.local`
  */
 export function loadEnvVariables(
   options?: EnvLoadOptions
 ): Record<string, string> | undefined {
   const envDir = options?.cwd ?? process.cwd();
-  const envFiles = DEFAULT_ENV_FILES.map((file) => path.join(envDir, file));
+  const envFiles = [
+    `.env`,
+    `.env.local`,
+    ...(options?.modes.flatMap((mode) => [
+      `.env.${mode}`,
+      `.env.${mode}.local`,
+    ]) ?? []),
+  ].map((file) => path.join(envDir, file));
 
   // Parse `envFiles` and combine their variables into a single object
   const parsed = Object.fromEntries(
diff --git a/libs/wingsdk/test/helpers.test.ts b/libs/wingsdk/test/helpers.test.ts
index ecfb8f8cc73..86d6d2dc9ef 100644
--- a/libs/wingsdk/test/helpers.test.ts
+++ b/libs/wingsdk/test/helpers.test.ts
@@ -13,6 +13,7 @@ describe("loadEnvVariables", () => {
 
     const loaded = loadEnvVariables({
       cwd: tempdir,
+      modes: ["test"],
     })!;
 
     expect(loaded).toBeDefined();
@@ -28,10 +29,35 @@ describe("loadEnvVariables", () => {
 
     const loaded = loadEnvVariables({
       cwd: tempdir,
+      modes: ["test"],
     })!;
 
     expect(loaded).toBeDefined();
     expect(loaded.BASE_THING).toBe("hi");
     expect(loaded.OTHER_THING).toBe("hi_wing");
   });
+
+  it("can load local and mode env files", async () => {
+    const tempdir = await mkdtemp(join(tmpdir(), "env-test"));
+
+    await writeFile(`${tempdir}/.env`, "TEST1=1\n");
+    await writeFile(`${tempdir}/.env.local`, "TEST2=2\n");
+    await writeFile(`${tempdir}/.env.run`, "TEST3=3\n");
+    await writeFile(`${tempdir}/.env.run.local`, "TEST4=4\n");
+    await writeFile(`${tempdir}/.env.it`, "TEST5=5\n");
+    await writeFile(`${tempdir}/.env.it.local`, "TEST6=6\n");
+
+    const loaded = loadEnvVariables({
+      cwd: tempdir,
+      modes: ["run", "it"],
+    })!;
+
+    expect(loaded).toBeDefined();
+    expect(loaded.TEST1).toBe("1");
+    expect(loaded.TEST2).toBe("2");
+    expect(loaded.TEST3).toBe("3");
+    expect(loaded.TEST4).toBe("4");
+    expect(loaded.TEST5).toBe("5");
+    expect(loaded.TEST6).toBe("6");
+  });
 });