A/B-ification of tests that can fail due to external influences

tests/conftest.py

@@ -1,8 +1,6 @@
 # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Pytest fixtures and redefined-outer-name don't mix well. Disable it.
-# pylint:disable=redefined-outer-name
 # We import some fixtures that are unused. Disable that too.
 # pylint:disable=unused-import
 

tests/framework/ab_test.py

@@ -23,6 +23,7 @@
 """
 import contextlib
 import os
+import shutil
 import statistics
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -31,6 +32,11 @@
 import scipy
 
 from framework import utils
+from framework.defs import FC_WORKSPACE_DIR
+from framework.microvm import Microvm
+from framework.utils import CommandReturn
+from framework.with_filelock import with_filelock
+from host_tools.cargo_build import get_firecracker_binaries
 
 # Locally, this will always compare against main, even if we try to merge into, say, a feature branch.
 # We might want to do a more sophisticated way to determine a "parent" branch here.
@@ -100,6 +106,133 @@ def git_ab_test(
         return result_a, result_b, comparison
 
 
+def is_pr() -> bool:
+    """Returns `True` iff we are executing in the context of a build kite run on a pull request"""
+    return os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false"
+
+
+def git_ab_test_host_command_if_pr(
+    command: str,
+    *,
+    comparator: Callable[[CommandReturn, CommandReturn], bool] = default_comparator,
+):
+    """Runs the given bash command as an A/B-Test if we're in a pull request context (asserting that its stdout and
+    stderr did not change across the PR). Otherwise runs the command, asserting it returns a zero exit code
+    """
+    if is_pr():
+        git_ab_test_host_command(command, comparator=comparator)
+    else:
+        utils.run_cmd(command)
+
+
+def git_ab_test_host_command(
+    command: str,
+    *,
+    comparator: Callable[[CommandReturn, CommandReturn], bool] = default_comparator,
+    a_revision: str = DEFAULT_A_REVISION,
+    b_revision: Optional[str] = None,
+):
+    """Performs an A/B-Test of the specified command, asserting that both the A and B invokations return the same stdout/stderr"""
+    (_, old_out, old_err), (_, new_out, new_err), the_same = git_ab_test(
+        lambda path, _is_a: utils.run_cmd(command, ignore_return_code=True, cwd=path),
+        comparator,
+        a_revision=a_revision,
+        b_revision=b_revision,
+    )
+
+    assert (
+        the_same
+    ), f"The output of running command `{command}` changed:\nOld:\nstdout:\n{old_out}\nstderr:\n{old_err}\n\nNew:\nstdout:\n{new_out}\nstderr:\n{new_err}"
+
+
+def set_did_not_grow_comparator(
+    set_generator: Callable[[CommandReturn], set]
+) -> Callable[[CommandReturn, CommandReturn], bool]:
+    """Factory function for comparators to use with git_ab_test_command that converts the command output to sets
+    (using the given callable) and then checks that the "B" set is a subset of the "A" set
+    """
+    return lambda output_a, output_b: set_generator(output_b).issubset(
+        set_generator(output_a)
+    )
+
+
+def git_ab_test_with_binaries(
+    test_runner: Callable[[Path, Path], T],
+    comparator: Callable[[T, T], U] = default_comparator,
+    *,
+    a_revision: str = DEFAULT_A_REVISION,
+    b_revision: Optional[str] = None,
+) -> (T, T, U):
+    """Similar to `git_ab_test`, with the only difference being that this function compiles firecracker at the specified
+    revisions and only passes the firecracker binaries to the test_runner. Maintains a cache of previously compiled
+    revisions, to prevent excessive recompilation across different tests of the same revision
+    """
+
+    @with_filelock
+    def grab_binaries(checkout: Path):
+        with chdir(checkout):
+            revision = utils.run_cmd("git rev-parse HEAD").stdout.strip()
+
+        revision_store = FC_WORKSPACE_DIR / "build" / revision
+        if not revision_store.exists():
+            firecracker, jailer = get_firecracker_binaries(workspace_dir=checkout)
+
+            revision_store.mkdir(parents=True, exist_ok=True)
+            shutil.copy(firecracker, revision_store / "firecracker")
+            shutil.copy(jailer, revision_store / "jailer")
+
+        return (
+            revision_store / "firecracker",
+            revision_store / "jailer",
+        )
+
+    return git_ab_test(
+        lambda checkout, _is_a: test_runner(*grab_binaries(checkout)),
+        comparator,
+        a_revision=a_revision,
+        b_revision=b_revision,
+    )
+
+
+def git_ab_test_guest_command(
+    microvm_factory: Callable[[Path, Path], Microvm],
+    command: str,
+    *,
+    comparator: Callable[[CommandReturn, CommandReturn], bool] = default_comparator,
+    a_revision: str = DEFAULT_A_REVISION,
+    b_revision: Optional[str] = None,
+):
+    """The same as git_ab_test_command, but via SSH. The closure argument should setup a microvm using the passed
+    paths to firecracker and jailer binaries."""
+
+    def test_runner(firecracker, jailer):
+        microvm = microvm_factory(firecracker, jailer)
+        return microvm.ssh.run(command)
+
+    (_, old_out, old_err), (_, new_out, new_err), the_same = git_ab_test_with_binaries(
+        test_runner, comparator, a_revision=a_revision, b_revision=b_revision
+    )
+
+    assert (
+        the_same
+    ), f"The output of running command `{command}` changed:\nOld:\nstdout:\n{old_out}\nstderr\n{old_err}\n\nNew:\nstdout:\n{new_out}\nstderr:\n{new_err}"
+
+
+def git_ab_test_guest_command_if_pr(
+    microvm_factory: Callable[[Path, Path], Microvm],
+    command: str,
+    *,
+    comparator=default_comparator,
+):
+    """The same as git_ab_test_command_if_pr, but via SSH"""
+    if is_pr():
+        git_ab_test_guest_command(microvm_factory, command, comparator=comparator)
+    else:
+        microvm = microvm_factory(*get_firecracker_binaries())
+        ecode, stdout, stderr = microvm.ssh.run(command)
+        assert ecode == 0, f"stdout:\n{stdout}\nstderr:\n{stderr}\n"
+
+
 def check_regression(
     a_samples: List[float], b_samples: List[float], *, n_resamples: int = 9999
 ):
@@ -147,6 +280,10 @@ def temporary_checkout(revision: str):
 
         yield Path(tmp_dir)
 
+        # If we compiled firecracker inside the checkout, python's recursive shutil.rmdir will
+        # run incredibly long. Thus, remove manually.
+        utils.run_cmd(f"rm -rf {tmp_dir}")
+
 
 # Once we upgrade to python 3.11, this will be in contextlib:
 # https://docs.python.org/3/library/contextlib.html#contextlib.chdir

tests/framework/defs.py

@@ -11,9 +11,6 @@
 # The Firecracker sources workspace dir
 FC_WORKSPACE_DIR = Path(__file__).parent.parent.parent.resolve()
 
-# Cargo target dir for the Firecracker workspace. Set via .cargo/config
-FC_WORKSPACE_TARGET_DIR = FC_WORKSPACE_DIR / "build/cargo_target"
-
 # Cargo build directory for seccompiler
 SECCOMPILER_TARGET_DIR = FC_WORKSPACE_DIR / "build/seccompiler"
 

tests/host_tools/cargo_build.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 
 from framework import defs, utils
-from framework.defs import FC_WORKSPACE_DIR, FC_WORKSPACE_TARGET_DIR
+from framework.defs import FC_WORKSPACE_DIR
 from framework.with_filelock import with_filelock
 
 CARGO_BUILD_REL_PATH = "firecracker_binaries"
@@ -66,32 +66,33 @@ def cargo_test(path, extra_args=""):
 
 
 @with_filelock
-def get_binary(name):
+def get_binary(name, *, workspace_dir=FC_WORKSPACE_DIR):
     """Build a binary"""
     target = DEFAULT_BUILD_TARGET
-    target_dir = FC_WORKSPACE_TARGET_DIR
-    out_dir = Path(f"{target_dir}/{target}/release")
-    bin_path = out_dir / name
+    target_dir = workspace_dir / "build" / "cargo_target"
+    bin_path = target_dir / target / "release" / name
     if not bin_path.exists():
         env = {"RUSTFLAGS": get_rustflags()}
         cargo(
             "build",
             f"-p {name} --release --target {target}",
             env=env,
-            cwd=FC_WORKSPACE_DIR,
+            cwd=workspace_dir,
         )
         utils.run_cmd(f"strip --strip-debug {bin_path}")
     return bin_path
 
 
 @with_filelock
-def get_firecracker_binaries():
+def get_firecracker_binaries(*, workspace_dir=FC_WORKSPACE_DIR):
     """Build the Firecracker and Jailer binaries if they don't exist.
 
     Returns the location of the firecracker related binaries eventually after
     building them in case they do not exist at the specified root_path.
     """
-    return get_binary("firecracker"), get_binary("jailer")
+    return get_binary("firecracker", workspace_dir=workspace_dir), get_binary(
+        "jailer", workspace_dir=workspace_dir
+    )
 
 
 @with_filelock

tests/host_tools/ip_generator.py

@@ -1,8 +1,6 @@
 # Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# pylint:disable=redefined-outer-name
-
 """
 Each microVM needs to have a unique IP on the host network, or there will be
 conflicts.

tests/integration_tests/functional/test_drives.py

@@ -2,8 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """Tests for guest-side operations on /drives resources."""
 
-# pylint:disable=redefined-outer-name
-
 import os
 from subprocess import check_output
 

tests/integration_tests/functional/test_rng.py

@@ -2,8 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """Tests for the virtio-rng device"""
 
-# pylint:disable=redefined-outer-name
-
 import pytest
 
 from framework.properties import global_props

tests/integration_tests/performance/conftest.py

@@ -1,9 +1,6 @@
 # Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Pytest fixtures and redefined-outer-name don't mix well. Disable it.
-# pylint:disable=redefined-outer-name
-
 """Fixtures for performance tests"""
 
 import json

tests/integration_tests/performance/test_boottime.py

@@ -2,8 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """Tests that ensure the boot time to init process is within spec."""
 
-# pylint:disable=redefined-outer-name
-
 import re
 import time
 

tests/integration_tests/performance/test_network_ab.py

@@ -67,9 +67,7 @@ def network_microvm(request, microvm_factory, guest_kernel, rootfs):
 
 @pytest.mark.nonci
 @pytest.mark.parametrize("network_microvm", [1], indirect=True)
-def test_network_latency(
-    network_microvm, metrics
-):  # pylint:disable=redefined-outer-name
+def test_network_latency(network_microvm, metrics):
     """
     Test network latency for multiple vm configurations.
 
@@ -134,7 +132,7 @@ def test_network_tcp_throughput(
     payload_length,
     mode,
     metrics,
-):  # pylint:disable=redefined-outer-name
+):
     """
     Iperf between guest and host in both directions for TCP workload.
     """

tests/integration_tests/performance/test_process_startup_time.py

@@ -2,8 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """Test that the process startup time up to socket bind is within spec."""
 
-# pylint: disable=redefined-outer-name
-
 import os
 import time
 

tests/integration_tests/security/test_jail.py

@@ -2,8 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """Tests that verify the jailer's behavior."""
 
-# pylint: disable=redefined-outer-name
-
 import functools
 import http.client as http_client
 import os

tests/integration_tests/security/test_sec_audit.py

@@ -1,26 +1,40 @@
 # Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Tests ensuring security vulnerabilities are not present in dependencies."""
-
+import json
 
 import pytest
 
-from framework import defs
+from framework.ab_test import (
+    git_ab_test_host_command_if_pr,
+    set_did_not_grow_comparator,
+)
+from framework.utils import CommandReturn
 from framework.utils_cpuid import CpuVendor, get_cpu_vendor
-from host_tools.cargo_build import cargo
 
 
 @pytest.mark.skipif(
     get_cpu_vendor() != CpuVendor.INTEL,
-    reason="The audit is based on cargo.lock which " "is identical on all platforms",
+    reason="The audit is based on cargo.lock which is identical on all platforms",
 )
 def test_cargo_audit():
     """
     Run cargo audit to check for crates with security vulnerabilities.
     """
-    # Run command and raise exception if non-zero return code
-    cargo(
-        "audit",
-        "--deny warnings -q",
-        cwd=defs.FC_WORKSPACE_DIR,
+
+    def set_of_vulnerabilities(output: CommandReturn):
+        output = json.loads(output.stdout)
+
+        return set(
+            frozenset(vulnerability)
+            for vulnerability in output["vulnerabilities"]["list"]
+        ).union(
+            frozenset(warning)
+            for warning_kind, warnings in output["warnings"].items()
+            for warning in warnings
+        )
+
+    git_ab_test_host_command_if_pr(
+        "cargo audit --deny warnings -q --json",
+        comparator=set_did_not_grow_comparator(set_of_vulnerabilities),
     )

tests/integration_tests/security/test_seccomp.py

@@ -90,7 +90,6 @@ def test_seccomp_ls(bin_seccomp_paths):
     """
     Assert that the seccomp filter denies an unallowed syscall.
     """
-    # pylint: disable=redefined-outer-name
     # pylint: disable=subprocess-run-check
     # The fixture pattern causes a pylint false positive for that rule.
 
@@ -135,7 +134,6 @@ def test_advanced_seccomp(bin_seccomp_paths):
     binary, denies the malicious demo binary and that an empty allowlist
     denies everything.
     """
-    # pylint: disable=redefined-outer-name
     # pylint: disable=subprocess-run-check
     # The fixture pattern causes a pylint false positive for that rule.
 
@@ -272,7 +270,6 @@ def test_seccomp_rust_panic(bin_seccomp_paths):
     Test that the Firecracker filters allow a Rust panic to run its
     course without triggering a seccomp violation.
     """
-    # pylint: disable=redefined-outer-name
     # pylint: disable=subprocess-run-check
     # The fixture pattern causes a pylint false positive for that rule.