buildkite/bazelci_test.py

#!/usr/bin/env python3
#
# Copyright 2023 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

os.environ["BUILDKITE_ORGANIZATION_SLUG"] = "bazel"
os.environ["BUILDKITE_PIPELINE_SLUG"] = "test"

import bazelci
import shlex
import tempfile
import unittest
import yaml


class CalculateFlags(unittest.TestCase):
    _CONFIGS = yaml.safe_load(
        """
.base_flags: &base_flags
  ? "--enable_a"
  ? "--enable_b"

tasks:
  basic:
    build_flags:
      - "--enable_x"
      - "--enable_y"
    build_targets:
      - "//..."
  json_profile:
    build_flags:
      - "--enable_x"
      - "--enable_y"
    build_targets:
      - "//..."
    include_json_profile:
      - build
  capture_corrupted:
    build_flags:
      - "--enable_x"
      - "--enable_y"
    build_targets:
      - "//..."
    capture_corrupted_outputs:
      - build
  no_flags:
    test_targets:
      - "//..."
  merge_flags:
    build_flags:
      <<: *base_flags
      ? "--enable_z"
      ? "--enable_w"
    build_targets:
      - "//..."
  flatten_flags:
    build_flags:
      - "--enable_z"
      - ["--enable_y", "--enable_x"]
      - "--enable_w"
    """
    )

    def test_basic_functionality(self):
        tasks = self._CONFIGS.get("tasks")
        flags, json_profile_out, capture_corrupted_outputs_dir = bazelci.calculate_flags(
            tasks.get("basic"), "build_flags", "build", "/tmp", ["HOME"]
        )
        self.assertEqual(flags, ["--enable_x", "--enable_y", "--test_env=HOME"])
        self.assertEqual(json_profile_out, None)
        self.assertEqual(capture_corrupted_outputs_dir, None)

    def test_json_profile(self):
        tasks = self._CONFIGS.get("tasks")
        flags, json_profile_out, capture_corrupted_outputs_dir = bazelci.calculate_flags(
            tasks.get("json_profile"), "build_flags", "build", "/tmp", ["HOME"]
        )
        self.assertEqual(
            flags,
            ["--enable_x", "--enable_y", "--profile=/tmp/build.profile.gz", "--test_env=HOME"],
        )
        self.assertEqual(json_profile_out, "/tmp/build.profile.gz")

    def test_capture_corrupted(self):
        tasks = self._CONFIGS.get("tasks")
        flags, json_profile_out, capture_corrupted_outputs_dir = bazelci.calculate_flags(
            tasks.get("capture_corrupted"), "build_flags", "build", "/tmp", ["HOME"]
        )
        self.assertEqual(
            flags,
            [
                "--enable_x",
                "--enable_y",
                "--experimental_remote_capture_corrupted_outputs=/tmp/build_corrupted_outputs",
                "--test_env=HOME",
            ],
        )
        self.assertEqual(capture_corrupted_outputs_dir, "/tmp/build_corrupted_outputs")

    def test_no_flags_in_config(self):
        tasks = self._CONFIGS.get("tasks")
        flags, json_profile_out, capture_corrupted_outputs_dir = bazelci.calculate_flags(
            tasks.get("no_flags"), "build_flags", "build", "/tmp", ["HOME"]
        )
        self.assertEqual(flags, ["--test_env=HOME"])

    def test_merge_flags(self):
        tasks = self._CONFIGS.get("tasks")
        flags, json_profile_out, capture_corrupted_outputs_dir = bazelci.calculate_flags(
            tasks.get("merge_flags"), "build_flags", "build", "/tmp", ["HOME"]
        )
        self.assertEqual(
            flags,
            ["--enable_a", "--enable_b", "--enable_z", "--enable_w", "--test_env=HOME"],
        )

    def test_flatten_flags(self):
        tasks = self._CONFIGS.get("tasks")
        flags, json_profile_out, capture_corrupted_outputs_dir = bazelci.calculate_flags(
            tasks.get("flatten_flags"), "build_flags", "build", "/tmp", ["HOME"]
        )
        self.assertEqual(
            flags,
            [
                "--enable_z",
                "--enable_y",
                "--enable_x",
                "--enable_w",
                "--test_env=HOME"],
        )


class CalculateTargets(unittest.TestCase):
    _CONFIGS = yaml.safe_load(
        """
.base_targets: &base_targets
  ? "//..."
  ? "-//experimental/..."

tasks:
  basic:
    build_targets:
      - "//..."
      - "-//bad/..."
  merge:
    build_targets:
      <<: *base_targets
      ? "//experimental/good/..."

  flatten:
    build_targets:
      - "//..."
      - ["-//bad/one", "-//bad/two"]
      - "-//bad/three"
    """
    )

    def test_basic_functionality(self):
        tasks = self._CONFIGS.get("tasks")
        build_targets, test_targets, coverage_targets, index_targets = bazelci.calculate_targets(
            tasks.get("basic"),
            "bazel",
            build_only=False,
            test_only=False,
            workspace_dir="/tmp",
            ws_setup_func=None,
            git_commit="abcd",
            test_flags=[],
        )
        self.assertEqual(build_targets, ["//...", "-//bad/..."])
        self.assertEqual(test_targets, [])
        self.assertEqual(coverage_targets, [])
        self.assertEqual(index_targets, [])

    def test_merge(self):
        tasks = self._CONFIGS.get("tasks")
        build_targets, test_targets, coverage_targets, index_targets = bazelci.calculate_targets(
            tasks.get("merge"),
            "bazel",
            build_only=False,
            test_only=False,
            workspace_dir="/tmp",
            ws_setup_func=None,
            git_commit="abcd",
            test_flags=[],
        )
        self.assertEqual(build_targets, ["//...", "-//experimental/...", "//experimental/good/..."])

    def test_flatten(self):
        tasks = self._CONFIGS.get("tasks")
        build_targets, test_targets, coverage_targets, index_targets = bazelci.calculate_targets(
            tasks.get("flatten"),
            "bazel",
            build_only=False,
            test_only=False,
            workspace_dir="/tmp",
            ws_setup_func=None,
            git_commit="abcd",
            test_flags=[],
        )
        self.assertEqual(build_targets, [
            "//...",
            "-//bad/one",
            "-//bad/two",
            "-//bad/three",
        ])

class MatrixExpansion(unittest.TestCase):
    _CONFIGS = yaml.safe_load(
        """
matrix:
  bazel: ["1.2.3", "2.3.4"]
  platform: ["pf1", "pf2"]
  exit_status: [1, 2]
tasks:
  basic:
    name: "Basic"
  unformatted:
    name: "Unformatted"
    bazel: ${{ bazel }}
  without_name:
    bazel: ${{ bazel }}
  formatted:
    name: "Formatted w/ Bazel v{bazel} on {platform}"
    bazel: ${{ bazel }}
    platform: ${{ platform }}
  exit_status:
    name: "Exit Status: {exit_status}"
    soft_fail:
      - exit_status: ${{exit_status}}
    """
    )

    def test_basic_functionality(self):
        import copy

        config = copy.deepcopy(self._CONFIGS)

        bazelci.expand_task_config(config)
        expanded_tasks = config["tasks"]

        self.assertEqual(
            list(expanded_tasks.values()),
            [
                # no matrix expansion
                dict(name="Basic"),
                dict(name="Unformatted", bazel="1.2.3"),
                dict(name="Unformatted", bazel="2.3.4"),
                # no name
                dict(bazel="1.2.3"),
                dict(bazel="2.3.4"),
                dict(name="Formatted w/ Bazel v1.2.3 on pf1", bazel="1.2.3", platform="pf1"),
                dict(name="Formatted w/ Bazel v1.2.3 on pf2", bazel="1.2.3", platform="pf2"),
                dict(name="Formatted w/ Bazel v2.3.4 on pf1", bazel="2.3.4", platform="pf1"),
                dict(name="Formatted w/ Bazel v2.3.4 on pf2", bazel="2.3.4", platform="pf2"),
                dict(name="Exit Status: 1", soft_fail=[dict(exit_status=1)]),
                dict(name="Exit Status: 2", soft_fail=[dict(exit_status=2)]),
            ],
        )

    def test_embedded_matrix_expansion(self):
        config = yaml.safe_load(
            """
matrix:
  bazel_version:
    7.x: 7.7.1
    8.x: 8.6.0
tasks:
  embedded:
    name: "Embedded {bazel_version}"
    build_flags:
      - --config=bazel-${{ bazel_version }}
      - --bazel_version=${{ bazel_version }}
      - prefix-${{ bazel_version }}-suffix
            """
        )

        bazelci.expand_task_config(config)
        expanded_tasks = list(config["tasks"].values())

        self.assertEqual(
            expanded_tasks,
            [
                dict(
                    name="Embedded 7.x",
                    build_flags=[
                        "--config=bazel-7.7.1",
                        "--bazel_version=7.7.1",
                        "prefix-7.7.1-suffix",
                    ],
                ),
                dict(
                    name="Embedded 8.x",
                    build_flags=[
                        "--config=bazel-8.6.0",
                        "--bazel_version=8.6.0",
                        "prefix-8.6.0-suffix",
                    ],
                ),
            ],
        )


class MatrixExclude(unittest.TestCase):
    _CONFIGS_SINGLE_EXCLUDE = yaml.safe_load(
        """
matrix:
  bazel: ["1.2.3", "2.3.4"]
  platform: ["pf1", "pf2"]
  exclude:
    - bazel: "1.2.3"
      platform: "pf2"
tasks:
  formatted:
    name: "Formatted w/ Bazel v{bazel} on {platform}"
    bazel: ${{ bazel }}
    platform: ${{ platform }}
        """
    )

    _CONFIGS_MULTIPLE_EXCLUDES = yaml.safe_load(
        """
matrix:
  bazel: ["1.2.3", "2.3.4"]
  platform: ["pf1", "pf2"]
  exclude:
    - bazel: "1.2.3"
      platform: "pf2"
    - bazel: "2.3.4"
      platform: "pf1"
tasks:
  formatted:
    name: "Formatted w/ Bazel v{bazel} on {platform}"
    bazel: ${{ bazel }}
    platform: ${{ platform }}
        """
    )

    _CONFIGS_PARTIAL_EXCLUDE = yaml.safe_load(
        """
matrix:
  bazel: ["1.2.3", "2.3.4"]
  platform: ["pf1", "pf2"]
  exclude:
    - platform: "pf2"
tasks:
  formatted:
    name: "Formatted w/ Bazel v{bazel} on {platform}"
    bazel: ${{ bazel }}
    platform: ${{ platform }}
        """
    )

    _CONFIGS_DICT_EXCLUDE = yaml.safe_load(
        """
matrix:
  flags:
    set-one: ["--one", "--two"]
    set-two: ["--three", "--four"]
  platform: ["pf1", "pf2"]
  exclude:
    - platform: "pf2"
      flags: set-one
tasks:
  formatted:
    name: "Formatted on {platform} with flags {flags}"
    platform: ${{ platform }}
    build_flags: ${{ flags }}
        """
    )

    def test_single_exclude(self):
        import copy
        config = copy.deepcopy(self._CONFIGS_SINGLE_EXCLUDE)

        bazelci.expand_task_config(config)
        expanded_tasks = config["tasks"]
        # Total combinations: 2 * 2 = 4, minus 1 excluded = 3
        self.assertEqual(len(expanded_tasks), 3)
        expanded_task_names = [task.get("name", None) for id, task in expanded_tasks.items()]
        self.assertEqual(expanded_task_names, [
            "Formatted w/ Bazel v1.2.3 on pf1",
            "Formatted w/ Bazel v2.3.4 on pf1",
            "Formatted w/ Bazel v2.3.4 on pf2",
        ])

    def test_multiple_excludes(self):
        import copy
        config = copy.deepcopy(self._CONFIGS_MULTIPLE_EXCLUDES)

        bazelci.expand_task_config(config)
        expanded_tasks = config["tasks"]
        # Total combinations: 2 * 2 = 4, minus 2 excluded = 2
        self.assertEqual(len(expanded_tasks), 2)
        expanded_task_names = [task.get("name", None) for id, task in expanded_tasks.items()]
        self.assertEqual(expanded_task_names, [
            "Formatted w/ Bazel v1.2.3 on pf1",
            "Formatted w/ Bazel v2.3.4 on pf2",
        ])

    def test_partial_attribute_exclude(self):
        import copy
        config = copy.deepcopy(self._CONFIGS_PARTIAL_EXCLUDE)

        bazelci.expand_task_config(config)
        expanded_tasks = config["tasks"]
        # Total combinations: 2 * 2 = 4, minus 2 excluded (all pf2) = 2
        self.assertEqual(len(expanded_tasks), 2)
        expanded_task_names = [task.get("name", None) for id, task in expanded_tasks.items()]
        self.assertEqual(expanded_task_names, [
            "Formatted w/ Bazel v1.2.3 on pf1",
            "Formatted w/ Bazel v2.3.4 on pf1",
        ])

    def test_dict_attribute_exclude(self):
        import copy
        config = copy.deepcopy(self._CONFIGS_DICT_EXCLUDE)

        bazelci.expand_task_config(config)
        expanded_tasks = config["tasks"]
        # Total combinations: 2 * 2 = 4, minus 1 excluded = 3
        self.assertEqual(len(expanded_tasks), 3)
        expanded_task_names = [task.get("name", None) for id, task in expanded_tasks.items()]
        # the name receives the friendly alias
        self.assertEqual(expanded_task_names, [
            "Formatted on pf1 with flags set-one",
            "Formatted on pf1 with flags set-two",
            "Formatted on pf2 with flags set-two",
        ])
        expanded_task_flags = [task.get("build_flags", None) for id, task in expanded_tasks.items()]
        # but the flags get the real value
        self.assertEqual(expanded_task_flags, [
            ["--one", "--two"],
            ["--three", "--four"],
            ["--three", "--four"],
        ])


class ShellQuoting(unittest.TestCase):
    """Regression tests: values interpolated into generated Buildkite step
    commands must be shell-quoted so that unusual task names (which come from
    user-controlled `.bazelci/presubmit.yml` dict keys) cannot break out of the
    --task argument and inject extra shell commands."""

    def _runner_command(self, step):
        # create_step stores the list of shell commands under "command".
        commands = step["command"]
        for cmd in commands:
            if "--task=" in cmd:
                return cmd
        self.fail("no command containing --task= was generated")

    def _sentinel_path(self):
        # A uniquely-named path that an unquoted injection would create as its
        # side effect. We never create it ourselves, so it must not exist after
        # the command is generated.
        path = tempfile.mktemp(prefix="pwned.")
        self.addCleanup(lambda: os.remove(path) if os.path.exists(path) else None)
        return path

    def test_runner_step_quotes_task_name(self):
        sentinel = self._sentinel_path()
        # If interpolated unquoted, the `;` would end the runner invocation and
        # `touch <sentinel>` would run as a separate command.
        task = "x; touch %s; #" % sentinel

        step = bazelci.runner_step(
            platform=bazelci.DEFAULT_PLATFORM,
            task=task,
            project_name="test",
        )
        cmd = self._runner_command(step)
        # Parse the command exactly as a shell would.
        tokens = shlex.split(cmd)
        # The whole task name must survive as a single --task= argument...
        self.assertIn("--task=" + task, tokens)
        # ...the injected `touch` must NOT appear as its own token...
        self.assertNotIn("touch", tokens)
        # ...and the injection's side-effect file must never be created.
        self.assertFalse(os.path.exists(sentinel))

    def test_bazel_build_step_quotes_task_name(self):
        sentinel = self._sentinel_path()
        task = "x; touch %s; #" % sentinel

        step = bazelci.bazel_build_step(
            task=task,
            platform=bazelci.DEFAULT_PLATFORM,
            project_name="test",
        )
        cmd = self._runner_command(step)
        tokens = shlex.split(cmd)
        self.assertIn("--task=" + task, tokens)
        self.assertNotIn("touch", tokens)
        self.assertFalse(os.path.exists(sentinel))

    def test_runner_step_quotes_config_arguments(self):
        sentinel = self._sentinel_path()

        step = bazelci.runner_step(
            platform=bazelci.DEFAULT_PLATFORM,
            task="ok",
            project_name="test",
            file_config="x; touch %s; #" % sentinel,
            git_commit="$(touch %s)" % sentinel,
        )
        cmd = self._runner_command(step)
        tokens = shlex.split(cmd)
        self.assertIn("--file_config=x; touch %s; #" % sentinel, tokens)
        self.assertIn("--git_commit=$(touch %s)" % sentinel, tokens)
        self.assertNotIn("touch", tokens)
        self.assertFalse(os.path.exists(sentinel))


if __name__ == "__main__":
    unittest.main()