[antlir2][image_test] support static listing

Summary:
Python and C++ tests in fbcode support "static listing" of test cases, which
means that in most cases TPX can list available test cases without actually
running the test binary for discovery.

Antlir and Antlir2 have never supported this, which is a huge disadvantage for
at least two reasons:

1. It's a huge perf hit when running tests on large booted images
2. Broken images can break test listing in ways that convince users it's
antlir's fault (especially booted tests with broken systemd units)

Test Plan:
```
❯ buck2 test fbcode//antlir/antlir2/testing/tests:test-cpp fbcode//antlir/antlir2/testing/tests:test-cpp-booted fbcode//antlir/antlir2/testing/tests:test-py fbcode//antlir/antlir2/testing/tests:test-py-booted
Buck UI: https://www.internalfb.com/buck2/269c3536-5543-4824-8ade-117d15d89fbd
Test UI: https://www.internalfb.com/intern/testinfra/testrun/2533275002169707
Note:    Using experimental modern dice
Network: Up: 170KiB  Down: 34KiB  (reSessionID-bdc336c4-eeed-4c44-a701-fc177d47d23a)
Jobs completed: 79. Time elapsed: 8.2s.
Cache hits: 69%. Commands: 16 (cached: 11, remote: 0, local: 5)
Tests finished: Pass 8. Fail 0. Fatal 0. Skip 0. Build failure 0
```

Reviewed By: epilatow

Differential Revision: D51271549

fbshipit-source-id: 6b1d1bbdb1e6eb3f41568185e04551fbfb24fe00

antlir/antlir2/test_images/install/BUCK

@@ -186,6 +186,7 @@ image_python_test(
     name = "test-installed-binary",
     srcs = ["test_installed_binary.py"],
     layer = ":installed-binary",
+    supports_static_listing = False,
 )
 
 export_file(

antlir/antlir2/testing/image_test.bzl

@@ -5,6 +5,7 @@
 
 # @oss-disable
 # @oss-disable
+load("//antlir/antlir2/bzl:macro_dep.bzl", "antlir2_dep")
 load("//antlir/antlir2/bzl:platform.bzl", "rule_with_default_target_platform")
 load("//antlir/antlir2/bzl:types.bzl", "LayerInfo")
 load("//antlir/antlir2/bzl/feature:defs.bzl", "feature")
@@ -76,13 +77,23 @@ def _impl(ctx: AnalysisContext) -> list[Provider]:
         is_executable = True,
         allow_args = True,
     )
+
+    env = dict(ctx.attrs.test[ExternalRunnerTestInfo].env)
+    if ctx.attrs._static_list_wrapper:
+        original = env.pop("STATIC_LIST_TESTS_BINARY", None)
+        if original:
+            env["STATIC_LIST_TESTS_BINARY"] = RunInfo(cmd_args(
+                ctx.attrs._static_list_wrapper[RunInfo],
+                cmd_args(original, format = "--wrap={}"),
+            ))
+
     return [
         ExternalRunnerTestInfo(
             command = [test_cmd],
             type = ctx.attrs.test[ExternalRunnerTestInfo].test_type,
             labels = ctx.attrs.labels + inner_labels,
             contacts = ctx.attrs.test[ExternalRunnerTestInfo].contacts,
-            env = ctx.attrs.test[ExternalRunnerTestInfo].env,
+            env = env,
             run_from_project_root = True,
         ),
         RunInfo(test_cmd),
@@ -131,6 +142,7 @@ _image_test = rule(
         "layer": attrs.dep(providers = [LayerInfo]),
         "run_as_user": attrs.string(default = "root"),
         "test": attrs.dep(providers = [ExternalRunnerTestInfo]),
+        "_static_list_wrapper": attrs.option(attrs.exec_dep(), default = None),
     } | cfg_attrs(),
     doc = "Run a test inside an image layer",
     cfg = layer_cfg,
@@ -156,6 +168,7 @@ def _implicit_image_test(
         _add_outer_labels: list[str] = [],
         default_os: str | None = None,
         # @oss-disable
+        _static_list_wrapper: str | None = None,
         visibility: list[str] | None = None,
         **kwargs):
     test_rule(
@@ -199,13 +212,14 @@ def _implicit_image_test(
         allocate_loop_devices = allocate_loop_devices,
         default_os = default_os,
         # @oss-disable
+        _static_list_wrapper = _static_list_wrapper,
         visibility = visibility,
     )
 
 image_cpp_test = partial(
     _implicit_image_test,
     cpp_unittest,
-    supports_static_listing = False,
+    _static_list_wrapper = antlir2_dep("testing/image_test:static-list-cpp"),
     _add_outer_labels = ["tpx:optout-test-result-output-spec"],
 )
 
@@ -238,8 +252,8 @@ def image_python_test(
         test_rule = python_unittest,
         name = name,
         layer = test_layer,
-        supports_static_listing = False,
         default_os = default_os,
         # @oss-disable
+        _static_list_wrapper = antlir2_dep("testing/image_test:static-list-py"),
         **kwargs
     )

antlir/antlir2/testing/image_test/BUCK

@@ -1,4 +1,4 @@
-load("//antlir/bzl:build_defs.bzl", "export_file", "rust_binary", "rust_library")
+load("//antlir/bzl:build_defs.bzl", "buck_command_alias", "export_file", "python_binary", "rust_binary", "rust_library")
 
 oncall("antlir")
 
@@ -39,3 +39,23 @@ export_file(
     name = "antlir2_image_test.service",
     visibility = ["PUBLIC"],
 )
+
+python_binary(
+    name = "static-list",
+    srcs = ["static_list.py"],
+    main_function = "antlir.antlir2.testing.image_test.static_list.main",
+)
+
+buck_command_alias(
+    name = "static-list-cpp",
+    args = ["cpp"],
+    exe = ":static-list",
+    visibility = ["PUBLIC"],
+)
+
+buck_command_alias(
+    name = "static-list-py",
+    args = ["py"],
+    exe = ":static-list",
+    visibility = ["PUBLIC"],
+)

antlir/antlir2/testing/image_test/src/main.rs

@@ -77,8 +77,8 @@ struct Boot {
     wants_units: Vec<String>,
 }
 
-#[derive(Parser, Debug)]
 /// Run a unit test inside an image layer.
+#[derive(Parser, Debug)]
 struct Args {
     #[clap(long)]
     spec: JsonFile<TestSpec>,
@@ -87,8 +87,6 @@ struct Args {
 }
 
 fn main() -> Result<()> {
-    let args = Args::parse();
-
     tracing_subscriber::registry()
         .with(
             tracing_subscriber::fmt::Layer::default()
@@ -102,6 +100,8 @@ fn main() -> Result<()> {
         .with(tracing_subscriber::EnvFilter::from_default_env())
         .init();
 
+    let args = Args::parse();
+
     let repo = find_root::find_repo_root(
         &absolute_path::AbsolutePathBuf::canonicalize(
             std::env::current_exe().context("while getting argv[0]")?,

antlir/antlir2/testing/image_test/static_list.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+The way that TPX invokes static list binaries is very unconducive to supporting
+it within antlir2.
+
+Since this requires some really hacky interpretation of argv, this is a simple
+python script that branches on the test type and then parses args and forwards
+them to the underlying static listing implementation.
+
+The args as received by this script are a terrible mix of the args that
+antlir2's image_test produces for internal use, the args for the static lister
+and other args that the static lister in TPX thinks are necessary for vanilla
+tests of this type.
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+
+
+def main():
+    sys.argv.pop(0)
+    which = sys.argv.pop(0)
+    if which == "cpp":
+        parser = argparse.ArgumentParser()
+        parser.add_argument("image_test_bin")
+        parser.add_argument("--wrap", required=True)
+        parser.add_argument("--spec", required=True)
+        parser.add_argument("test_type", choices=["gtest"])
+        parser.add_argument("cmd", nargs="+")
+        args = parser.parse_args(sys.argv)
+        os.execv(args.wrap, [args.wrap] + args.cmd)
+    if which == "py":
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--wrap", required=True)
+        parser.add_argument("--spec", required=True)
+        parser.add_argument("--json-output", required=True)
+        parser.add_argument("image_test_bin")
+        parser.add_argument("test_type", choices=["pyunit"])
+        parser.add_argument("cmd", nargs="+")
+        args = parser.parse_args(sys.argv)
+        os.execv(args.wrap, [args.wrap, "--json-output", args.json_output] + args.cmd)
+
+    raise Exception(f"Unknown static list wrapper: {which}")
+
+
+if __name__ == "__main__":
+    main()