From 1647e31799442c5245b4d46e965787bae451170f Mon Sep 17 00:00:00 2001
From: Paul Gschwendtner <paulgschwendtner@gmail.com>
Date: Fri, 15 Oct 2021 19:23:43 +0200
Subject: [PATCH] feat(bazel): add support for custom environment variables in
 integration tests

Adds support for custom environment variables in integration tests. The
environment variable dictionary can use location expansion. Also there
is a special placeholder called `<TMP>` that can be used to acquire a
temporary directory. This is useful when setting up tools like BAZELISK,
which requires a `HOME` environment variable as an example.
---
 bazel/integration/index.bzl                   | 43 +++++++++-
 bazel/integration/test_runner/bazel.ts        | 12 +++
 bazel/integration/test_runner/constants.ts    | 14 ++++
 bazel/integration/test_runner/main.ts         |  8 +-
 bazel/integration/test_runner/runner.ts       | 78 +++++++++++++------
 .../tests/custom_env_variables/BUILD.bazel    | 22 ++++++
 .../tests/custom_env_variables/test.js        | 39 ++++++++++
 7 files changed, 190 insertions(+), 26 deletions(-)
 create mode 100644 bazel/integration/test_runner/constants.ts
 create mode 100644 bazel/integration/tests/custom_env_variables/BUILD.bazel
 create mode 100644 bazel/integration/tests/custom_env_variables/test.js

diff --git a/bazel/integration/index.bzl b/bazel/integration/index.bzl
index 9910be90a1..29fb86e746 100644
--- a/bazel/integration/index.bzl
+++ b/bazel/integration/index.bzl
@@ -6,10 +6,35 @@ def _serialize_file(file):
 
     return struct(path = file.path, shortPath = file.short_path)
 
+def _serialize_and_expand_location(ctx, value):
+    """Expands Bazel make location expressions for the given value. Returns a JSON
+      serializable dictionary matching the `BazelExpandedValue` type in the test runner."""
+    new_value = ctx.expand_location(value, targets = ctx.attr.data)
+
+    return {
+        "value": new_value,
+        "containsExpandedValue": new_value != value,
+    }
+
 def _split_and_expand_command(ctx, command):
     """Splits a command into the binary and its arguments. Also Bazel locations are expanded."""
-    expanded_command = ctx.expand_location(command, targets = ctx.attr.data)
-    return expanded_command.split(" ", 1)
+    command, args = command.split(" ", 1)
+
+    return [
+        _serialize_and_expand_location(ctx, command),
+        _serialize_and_expand_location(ctx, args),
+    ]
+
+def _serialize_and_expand_environment(ctx, environment_dict):
+    """Converts the given environment dictionary into a JSON-serializable dictionary
+      that will work with the test runner."""
+    result = {}
+
+    for variable_name in environment_dict:
+        value = environment_dict[variable_name]
+        result[variable_name] = _serialize_and_expand_location(ctx, value)
+
+    return result
 
 def _unwrap_label_keyed_mappings(dict, description):
     """Unwraps a label-keyed dictionary used for expressing mappings into a JSON-serializable
@@ -51,6 +76,7 @@ def _integration_test_config_impl(ctx):
         testPackage = ctx.label.package,
         testFiles = [_serialize_file(f) for f in ctx.files.srcs],
         commands = [_split_and_expand_command(ctx, c) for c in ctx.attr.commands],
+        environment = _serialize_and_expand_environment(ctx, ctx.attr.environment),
         npmPackageMappings = npmPackageMappings,
         toolMappings = toolMappings,
     )
@@ -105,6 +131,17 @@ _integration_test_config = rule(
               This allows for binaries like `node` to be made available to the integration test
               using the `PATH` environment variable.""",
         ),
+        "environment": attr.string_dict(
+            doc = """
+              Dictionary of environment variables and their values. This allows for custom
+              environment variables to be set when integration commands are invoked.
+
+              The environment variable values can use Bazel make location expansion similar
+              to the `commands` attribute. Additionally, values of `<TMP>` are replaced with
+              a unique temporary directory. This can be useful when providing `HOME` for
+              bazelisk or puppeteer as as an example.
+            """,
+        ),
     },
 )
 
@@ -114,6 +151,7 @@ def integration_test(
         commands,
         npm_packages = {},
         tool_mappings = {},
+        environment = {},
         data = [],
         tags = [],
         **kwargs):
@@ -129,6 +167,7 @@ def integration_test(
         commands = commands,
         npm_packages = npm_packages,
         tool_mappings = tool_mappings,
+        environment = environment,
         tags = tags,
     )
 
diff --git a/bazel/integration/test_runner/bazel.ts b/bazel/integration/test_runner/bazel.ts
index 0673754da1..e83c923366 100644
--- a/bazel/integration/test_runner/bazel.ts
+++ b/bazel/integration/test_runner/bazel.ts
@@ -23,6 +23,18 @@ export interface BazelFileInfo {
   shortPath: string;
 }
 
+/**
+ * Interface describing a Bazel-expanded value. A integration command for example could
+ * use a Bazel location expansion to resolve a binary. Such resolved values are captured in
+ * a structure like this.
+ */
+export interface BazelExpandedValue {
+  /** Actual value, with expanded Make expressions if it contained any. */
+  value: string;
+  /** Whether the value contains an expanded value. */
+  containsExpandedValue: boolean;
+}
+
 /** Resolves the specified Bazel file to an absolute disk path. */
 export function resolveBazelFile(file: BazelFileInfo): string {
   return runfiles.resolveWorkspaceRelative(file.shortPath);
diff --git a/bazel/integration/test_runner/constants.ts b/bazel/integration/test_runner/constants.ts
new file mode 100644
index 0000000000..01ad6c2e58
--- /dev/null
+++ b/bazel/integration/test_runner/constants.ts
@@ -0,0 +1,14 @@
+/**
+ * @license
+ * Copyright Google LLC All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.io/license
+ */
+
+/**
+ * Placeholder that can be used in the test environment dictionary to reserve a
+ * temporary directory. A temporary directory can be useful when running tools
+ * like Bazelisk which need a `HOME` directory for example.
+ */
+export const ENVIRONMENT_TMP_PLACEHOLDER = '<TMP>';
diff --git a/bazel/integration/test_runner/main.ts b/bazel/integration/test_runner/main.ts
index e7d042e7d8..354a3ccd57 100644
--- a/bazel/integration/test_runner/main.ts
+++ b/bazel/integration/test_runner/main.ts
@@ -7,7 +7,7 @@
  */
 
 import {TestRunner} from './runner';
-import {BazelFileInfo, runfiles} from './bazel';
+import {BazelExpandedValue, BazelFileInfo, runfiles} from './bazel';
 import * as fs from 'fs';
 import {debug} from './debug';
 
@@ -20,7 +20,8 @@ interface TestConfig {
   testFiles: BazelFileInfo[];
   npmPackageMappings: Record<string, BazelFileInfo>;
   toolMappings: Record<string, BazelFileInfo>;
-  commands: [[binary: string, ...args: string[]]];
+  commands: [[binary: BazelExpandedValue, ...args: BazelExpandedValue[]]];
+  environment: Record<string, BazelExpandedValue>;
 }
 
 /** Main command line entry-point for the integration test runner. */
@@ -32,12 +33,15 @@ async function main(): Promise<void> {
   const configContent = await fs.promises.readFile(configPath, 'utf8');
   const config = JSON.parse(configContent) as TestConfig;
 
+  debug('Fetched test config:', config);
+
   const runner = new TestRunner(
     config.testFiles,
     config.testPackage,
     config.toolMappings,
     config.npmPackageMappings,
     config.commands,
+    config.environment,
   );
 
   await runner.run();
diff --git a/bazel/integration/test_runner/runner.ts b/bazel/integration/test_runner/runner.ts
index ed082a50ac..63a0adc775 100644
--- a/bazel/integration/test_runner/runner.ts
+++ b/bazel/integration/test_runner/runner.ts
@@ -15,7 +15,12 @@ import {
   readPackageJsonContents,
   updateMappingsForPackageJson,
 } from './package_json';
-import {BazelFileInfo, resolveBazelFile, resolveBinaryWithRunfiles} from './bazel';
+import {
+  BazelExpandedValue,
+  BazelFileInfo,
+  resolveBazelFile,
+  resolveBinaryWithRunfiles,
+} from './bazel';
 import {debug} from './debug';
 import {
   expandEnvironmentVariableSubstitutions,
@@ -23,6 +28,7 @@ import {
   prependToPathVariable,
   runCommandInChildProcess,
 } from './process_utils';
+import {ENVIRONMENT_TMP_PLACEHOLDER} from './constants';
 
 /**
  * Test runner that takes a set of files within a Bazel package and copies the files
@@ -38,18 +44,20 @@ export class TestRunner {
     private readonly testPackage: string,
     private readonly toolMappings: Record<string, BazelFileInfo>,
     private readonly npmPackageMappings: Record<string, BazelFileInfo>,
-    private readonly commands: [[binary: string, ...args: string[]]],
+    private readonly commands: [[binary: BazelExpandedValue, ...args: BazelExpandedValue[]]],
+    private readonly environment: Record<string, BazelExpandedValue>,
   ) {}
 
   async run() {
-    const tmpDir = await this._getTmpDirectoryPath();
-    const toolMappings = await this._setupToolMappingsForTest(tmpDir);
+    const testDir = await this._getTestTmpDirectoryPath();
+    const toolMappings = await this._setupToolMappingsForTest(testDir);
+    const testEnv = await this._createTestProcessEnvironment(testDir, toolMappings.binDir);
 
-    console.info(`Running test in: ${path.normalize(tmpDir)}`);
+    console.info(`Running test in: ${path.normalize(testDir)}`);
 
-    await this._copyTestFilesToDirectory(tmpDir);
-    await this._patchPackageJsonIfNeeded(tmpDir);
-    await this._runTestCommands(tmpDir, toolMappings.binDir);
+    await this._copyTestFilesToDirectory(testDir);
+    await this._patchPackageJsonIfNeeded(testDir);
+    await this._runTestCommands(testDir, testEnv);
   }
 
   /**
@@ -59,7 +67,7 @@ export class TestRunner {
    * In case this test does not run as part of `bazel test`, a system-temporary directory
    * is being created, although not being cleaned up to allow for debugging.
    */
-  private async _getTmpDirectoryPath(): Promise<string> {
+  private async _getTestTmpDirectoryPath(): Promise<string> {
     // Bazel provides a temporary test directory itself when it executes a test. We prefer
     // this when the integration test runs with `bazel test`. In other cases we want to
     // provide a temporary directory that can be used for manually jumping into the
@@ -167,24 +175,50 @@ export class TestRunner {
   }
 
   /**
-   * Runs the test commands sequentially in the test directory. An additional directory
-   * that is added to the command process `$PATH` environment variables can be specified.
+   * Creates the test process environment. The specified tools bin directory
+   * will be added to the environment `PATH` variable.
    *
-   * @throws An error if any of the configured commands did not complete successfully.
+   * User-specified environment variables can use Bazel location expansion as well
+   * as a special placeholder for acquiring a temporary directory for the test.
    */
-  private async _runTestCommands(
+  private async _createTestProcessEnvironment(
     testDir: string,
-    additionalPathDirectory: string | null,
-  ): Promise<void> {
-    const commandPath =
-      additionalPathDirectory === null
-        ? process.env.PATH
-        : prependToPathVariable(additionalPathDirectory, process.env.PATH ?? '');
-    const commandEnv = {...process.env, PATH: commandPath};
+    toolsBinDir: string,
+  ): Promise<NodeJS.ProcessEnv> {
+    const testEnv: NodeJS.ProcessEnv = {...process.env};
+    let i = 0;
+
+    for (let [variableName, value] of Object.entries(this.environment)) {
+      let envValue: string = value.value;
+
+      if (value.containsExpandedValue) {
+        envValue = await resolveBinaryWithRunfiles(envValue);
+      } else if (envValue === ENVIRONMENT_TMP_PLACEHOLDER) {
+        envValue = path.join(testDir, `.tmp-env-${i++}`);
+        await fs.promises.mkdir(envValue);
+      }
 
+      testEnv[variableName] = envValue;
+    }
+
+    const commandPath = prependToPathVariable(toolsBinDir, testEnv.PATH ?? '');
+    return {...testEnv, PATH: commandPath};
+  }
+
+  /**
+   * Runs the test commands sequentially in the test directory with the given test
+   * environment applied to the child process executing the command.
+   *
+   * @throws An error if any of the configured commands did not complete successfully.
+   */
+  private async _runTestCommands(testDir: string, commandEnv: NodeJS.ProcessEnv): Promise<void> {
     for (const [binary, ...args] of this.commands) {
-      const resolvedBinary = await resolveBinaryWithRunfiles(binary);
-      const evaluatedArgs = expandEnvironmentVariableSubstitutions(args);
+      // Only resolve the binary if it contains an expanded value. In other cases we would
+      // not want to resolve through runfiles to avoid accidentally unexpected resolution.
+      const resolvedBinary = binary.containsExpandedValue
+        ? await resolveBinaryWithRunfiles(binary.value)
+        : binary.value;
+      const evaluatedArgs = expandEnvironmentVariableSubstitutions(args.map((v) => v.value));
       const success = await runCommandInChildProcess(
         resolvedBinary,
         evaluatedArgs,
diff --git a/bazel/integration/tests/custom_env_variables/BUILD.bazel b/bazel/integration/tests/custom_env_variables/BUILD.bazel
new file mode 100644
index 0000000000..ddce66bcba
--- /dev/null
+++ b/bazel/integration/tests/custom_env_variables/BUILD.bazel
@@ -0,0 +1,22 @@
+load("//bazel/integration:index.bzl", "integration_test")
+
+integration_test(
+    name = "test",
+    srcs = [
+        "test.js",
+    ],
+    commands = [
+        "node ./test.js",
+    ],
+    data = ["@npm//:node_modules/semver/package.json"],
+    environment = {
+        "CUSTOM_VAR": "yes!",
+        "RESOLVED_BIN": "$(rootpath @npm//:node_modules/semver/package.json)",
+        "MANIFEST_PATH": "external/npm/node_modules/semver/package.json",
+        "BAZELISK_HOME": "<TMP>",
+        "BAZELISK_HOME_2": "<TMP>",
+    },
+    tool_mappings = {
+        "@nodejs//:node_bin": "node",
+    },
+)
diff --git a/bazel/integration/tests/custom_env_variables/test.js b/bazel/integration/tests/custom_env_variables/test.js
new file mode 100644
index 0000000000..7df2d600da
--- /dev/null
+++ b/bazel/integration/tests/custom_env_variables/test.js
@@ -0,0 +1,39 @@
+const fs = require('fs');
+
+if (process.env.CUSTOM_VAR !== 'yes!') {
+  console.error('The expected `CUSTOM_VAR` environment variable is not set.');
+  process.exit(1);
+}
+
+if (process.env.RESOLVED_BIN === undefined) {
+  console.error('The expected `RESOLVED_BIN` variable is not set.');
+  process.exit(1);
+}
+
+if (require(process.env.RESOLVED_BIN).name !== 'semver') {
+  console.error('The `RESOLVED_BIN` file did not resolve to the "package.json" of "semver".');
+  process.exit(1);
+}
+
+if (process.env.MANIFEST_PATH !== 'external/npm/node_modules/semver/package.json') {
+  console.error('Expected `MANIFEST_PATH` to remain untouched as it has not not been expanded.');
+  process.exit(1);
+}
+
+const bazeliskHome = process.env.BAZELISK_HOME;
+const bazeliskHome_2 = process.env.BAZELISK_HOME_2;
+
+if (!fs.statSync(bazeliskHome).isDirectory()) {
+  console.error('Expected `BAZELISK_HOME` environment variable to point to a temp directory.');
+  process.exit(1);
+}
+
+if (!fs.statSync(bazeliskHome_2).isDirectory()) {
+  console.error('Expected `BAZELISK_HOME_2` environment variable to point to a temp directory.');
+  process.exit(1);
+}
+
+if (bazeliskHome === bazeliskHome_2) {
+  console.error('Expected the bazelisk home variables to point to different temp directories.');
+  process.exit(1);
+}