From 02b71d563b5e8075989b852a75168fe79bdd8927 Mon Sep 17 00:00:00 2001
From: Junjie Wu <junjiew@meta.com>
Date: Wed, 20 Sep 2023 11:21:55 -0700
Subject: [PATCH] [antlir2][vm] add missing label for python tests

Summary:
Some tests aren't happy with our existing labels for python tests,
even though it worked fine for some. Copy more magic from
antlir/vm/bzl/defs.bzl for the labels.
```
$ buck targets resctl/atropos/vm_tests:
File changed: fbcode//antlir/antlir2/antlir2_vm/bzl/test.bzl
Error parsing fbcode//resctl/atropos/vm_tests
Error evaluating build file: `fbcode//resctl/atropos/vm_tests:TARGETS`

Caused by:
    Traceback (most recent call last):
      * fbcode/resctl/atropos/vm_tests/TARGETS:25, in <module>
          vm.python_test(
      * fbcode/buck2/starlark-rust/starlark/src/stdlib/partial.rs:175, in partial
          <native>
      * fbcode/antlir/antlir2/antlir2_vm/bzl/test.bzl:141, in _implicit_vm_test
          test_rule(
      * fbcode/antlir/bzl/build_defs.bzl:174, in python_unittest
          shim.python_unittest(**kwargs)
      * fbcode/antlir/bzl/build_defs_impl.bzl:404, in _python_unittest
          _wrap_internal(python_unittest, args, kwargs)
      * fbcode/antlir/bzl/build_defs_impl.bzl:298, in _wrap_internal
          fn(*args, **kwargs)
      * tools/build_defs/fbcode_macros/build_defs/python_unittest.bzl:257, in python_unittest
          fail(
    error: fail: You are trying to run a new TestPilot adapter interface module but is not using the tag 'use-testpilot-adapter' to enable that.
       --> tools/build_defs/fbcode_macros/build_defs/python_unittest.bzl:257:13
        |
    257 |               fail(
        |  _____________^
    258 | |                 "You are trying to run a new TestPilot adapter interface module but " +
    259 | |                 "is not using the tag 'use-testpilot-adapter' to enable that.",
    260 | |             )
        | |_____________^
        |

Failed to parse 1 package
```

Test Plan: CI

Reviewed By: vmagro

Differential Revision: D49440968

fbshipit-source-id: a124d53bd10b1ca292cdb7b05576262e44d7c50e
---
 antlir/antlir2/antlir2_vm/bzl/test.bzl | 49 +++++++++++++++++++-------
 1 file changed, 37 insertions(+), 12 deletions(-)

diff --git a/antlir/antlir2/antlir2_vm/bzl/test.bzl b/antlir/antlir2/antlir2_vm/bzl/test.bzl
index 61123116fd..920bdb967c 100644
--- a/antlir/antlir2/antlir2_vm/bzl/test.bzl
+++ b/antlir/antlir2/antlir2_vm/bzl/test.bzl
@@ -115,6 +115,32 @@ _vm_test = rule(
 
 vm_test = rule_with_default_target_platform(_vm_test)
 
+def _get_internal_labels(test_rule, run_as_bundle: bool):
+    wrapper_labels = ["heavyweight"]
+    if run_as_bundle:
+        wrapper_labels.append(special_tags.run_as_bundle)
+    if fully_qualified_test_name_rollout.use_fully_qualified_name():
+        wrapper_labels = wrapper_labels + [NAMING_ROLLOUT_LABEL]
+    wrapper_labels.append(special_tags.enable_artifact_reporting)
+
+    inner_labels = add_test_framework_label(HIDE_TEST_LABELS, "test-framework=8:vmtest")
+
+    # Due to a complex internal migration, these labels are required to both
+    # change the runtime behavior of the outer test, as well as build-time
+    # behavior of the inner target.
+    if test_rule == python_unittest:
+        wrapper_labels.append("use-testpilot-adapter")
+        inner_labels.append("use-testpilot-adapter")
+
+        # this tag gets added to the inner test automatically, but we must
+        # inform tpx that the wrapper observes the same behavior
+        wrapper_labels.append("tpx:list-format-migration:json")
+
+        # also annotate wrapper target with a framework
+        wrapper_labels = add_test_framework_label(wrapper_labels, "test-framework=8:vmtest")
+
+    return inner_labels, wrapper_labels
+
 def _implicit_vm_test(
         test_rule,
         name: str,
@@ -137,27 +163,26 @@ def _implicit_vm_test(
         should not be undone by the test fixture (ie, rebooting or setting
         a sysctl that cannot be undone for example).
     """
+    wrapper_labels = list(labels) if labels else []
+    wrapper_labels.extend(_add_outer_labels)
+    inner_labels = []
+
+    # @oss-disable
+    # @oss-disable
+    # @oss-disable
+
     inner_test_name = name + "_vm_test_inner"
     test_rule(
         name = inner_test_name,
         antlir_rule = "user-internal",
-        labels = add_test_framework_label(HIDE_TEST_LABELS, "test-framework=8:vmtest"),
+        labels = inner_labels,
         **kwargs
     )
-    labels = list(labels) if labels else []
-    labels.append("heavyweight")
-    labels.extend(_add_outer_labels)
-
-    if run_as_bundle:
-        labels.append("run_as_bundle")
-
-    # @oss-disable
-        # @oss-disable
 
     vm_test(
         name = name,
         test = ":" + inner_test_name,
-        test_labels = labels + [special_tags.enable_artifact_reporting],
+        test_labels = wrapper_labels,
         vm_host = vm_host,
         timeout_secs = timeout_secs,
         # VM is not ready for other arch yet
@@ -167,7 +192,7 @@ def _implicit_vm_test(
 vm_cpp_test = partial(
     _implicit_vm_test,
     cpp_unittest,
-    _add_outer_labels = ["tpx:optout-test-result-output-spec"],
+    # @oss-disable
     supports_static_listing = False,
 )
 vm_python_test = partial(_implicit_vm_test, python_unittest, supports_static_listing = False)