Add build mode to test action keys

Reviewed By: krallin

Differential Revision: D48970919

fbshipit-source-id: 6f3ce14507c300d77f1db70f13a08744ab1b57be

Author: c-ryan747

prelude/build_mode.bzl

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under both the MIT license found in the
+# LICENSE-MIT file in the root directory of this source tree and the Apache
+# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+# of this source tree.
+
+BuildModeInfo = provider(
+    fields = ["cell", "mode"],
+)

prelude/cxx/cxx.bzl

@@ -620,7 +620,7 @@ def cxx_test_impl(ctx: AnalysisContext) -> list[Provider]:
     command = [cmd_args(output.binary).hidden(output.runtime_files)] + ctx.attrs.args
 
     # Setup a RE executor based on the `remote_execution` param.
-    re_executor = get_re_executor_from_props(ctx.attrs.remote_execution)
+    re_executor = get_re_executor_from_props(ctx)
 
     return inject_test_run_info(
         ctx,

prelude/decls/common.bzl

@@ -12,6 +12,8 @@
 # the generated docs, and so those should be verified to be accurate and
 # well-formatted (and then delete this TODO)
 
+load("@prelude//:build_mode.bzl", "BuildModeInfo")
+
 def validate_uri(_s):
     return True
 
@@ -234,6 +236,15 @@ def _re_opts_for_tests_arg() -> "attribute":
         default = None,
     )
 
+def _re_action_key_provider_arg() -> "attribute":
+    return attrs.dep(providers = [BuildModeInfo], default = "fbcode//buck2/platform/build_mode:build_mode")
+
+def _re_test_args() -> dict[str, "attribute"]:
+    return {
+        "remote_execution": _re_opts_for_tests_arg(),
+        "remote_execution_action_key_providers": _re_action_key_provider_arg(),
+    }
+
 def _allow_cache_upload_arg():
     return {
         "allow_cache_upload": attrs.bool(
@@ -261,7 +272,7 @@ buck = struct(
     run_test_separately_arg = _run_test_separately_arg,
     fork_mode = _fork_mode,
     test_rule_timeout_ms = _test_rule_timeout_ms,
-    re_opts_for_tests_arg = _re_opts_for_tests_arg,
+    re_test_args = _re_test_args,
     target_os_type_arg = _target_os_type_arg,
     allow_cache_upload_arg = _allow_cache_upload_arg,
 )

prelude/decls/rust_rules.bzl

@@ -326,8 +326,8 @@ rust_test = prelude_rule(
         _RUST_EXECUTABLE_ATTRIBUTES |
         {
             "framework": attrs.bool(default = True),
-            "remote_execution": buck.re_opts_for_tests_arg(),
         } |
+        buck.re_test_args() |
         rust_common.toolchains_args() |
         rust_common.workspaces_arg()
     ),

prelude/decls/shell_rules.bzl

@@ -181,13 +181,12 @@ sh_test = prelude_rule(
             "licenses": attrs.list(attrs.source(), default = []),
             "list_args": attrs.option(attrs.list(attrs.string()), default = None),
             "list_env": attrs.option(attrs.dict(key = attrs.string(), value = attrs.string(), sorted = False), default = None),
-            "remote_execution": buck.re_opts_for_tests_arg(),
             "resources": attrs.list(attrs.source(), default = []),
             "run_args": attrs.list(attrs.string(), default = []),
             "run_env": attrs.dict(key = attrs.string(), value = attrs.string(), sorted = False, default = {}),
             "run_test_separately": attrs.bool(default = False),
             "test_rule_timeout_ms": attrs.option(attrs.int(), default = None),
-        }
+        } | buck.re_test_args()
     ),
 )
 

prelude/python/python_needed_coverage_test.bzl

@@ -43,7 +43,7 @@ def python_needed_coverage_test_impl(ctx: AnalysisContext) -> list[Provider]:
     test_env["TEST_PILOT"] = "1"
 
     # Setup a RE executor based on the `remote_execution` param.
-    re_executor = get_re_executor_from_props(ctx.attrs.remote_execution)
+    re_executor = get_re_executor_from_props(ctx)
 
     return inject_test_run_info(
         ctx,

prelude/python/python_test.bzl

@@ -62,7 +62,7 @@ def python_test_impl(ctx: AnalysisContext) -> list[Provider]:
     test_cmd = pex.run_cmd
 
     # Setup a RE executor based on the `remote_execution` param.
-    re_executor = get_re_executor_from_props(ctx.attrs.remote_execution)
+    re_executor = get_re_executor_from_props(ctx)
 
     return inject_test_run_info(
         ctx,

prelude/rules_impl.bzl

@@ -348,10 +348,8 @@ def _python_executable_attrs():
 
 def _python_test_attrs():
     test_attrs = _python_executable_attrs()
-    test_attrs.update({
-        "remote_execution": buck.re_opts_for_tests_arg(),
-        "_test_main": attrs.source(default = "prelude//python/tools:__test_main__.py"),
-    })
+    test_attrs["_test_main"] = attrs.source(default = "prelude//python/tools:__test_main__.py")
+    test_attrs.update(buck.re_test_args())
     return test_attrs
 
 def _cxx_binary_and_test_attrs():
@@ -435,10 +433,7 @@ inlined_extra_attributes = {
         "_omnibus_environment": omnibus_environment_attr(),
     },
     "cxx_python_extension": _cxx_python_extension_attrs(),
-    "cxx_test": dict(
-        remote_execution = buck.re_opts_for_tests_arg(),
-        **_cxx_binary_and_test_attrs()
-    ),
+    "cxx_test": buck.re_test_args() | _cxx_binary_and_test_attrs(),
     "cxx_toolchain": cxx_toolchain_extra_attributes(is_toolchain_rule = False),
 
     # Generic rule to build from a command
@@ -550,14 +545,14 @@ inlined_extra_attributes = {
         "_cxx_toolchain": toolchains_common.cxx(),
         "_python_toolchain": toolchains_common.python(),
     },
-    "python_needed_coverage_test": {
-        "contacts": attrs.list(attrs.string(), default = []),
-        "env": attrs.dict(key = attrs.string(), value = attrs.arg(), sorted = False, default = {}),
-        "labels": attrs.list(attrs.string(), default = []),
-        "needed_coverage": attrs.list(attrs.tuple(attrs.int(), attrs.dep(), attrs.option(attrs.string())), default = []),
-        "remote_execution": buck.re_opts_for_tests_arg(),
-        "test": attrs.dep(providers = [ExternalRunnerTestInfo]),
-    },
+    "python_needed_coverage_test": dict(
+        contacts = attrs.list(attrs.string(), default = []),
+        env = attrs.dict(key = attrs.string(), value = attrs.arg(), sorted = False, default = {}),
+        labels = attrs.list(attrs.string(), default = []),
+        needed_coverage = attrs.list(attrs.tuple(attrs.int(), attrs.dep(), attrs.option(attrs.string())), default = []),
+        test = attrs.dep(providers = [ExternalRunnerTestInfo]),
+        **buck.re_test_args()
+    ),
     "python_test": _python_test_attrs(),
     "remote_file": {
         "sha1": attrs.option(attrs.string(), default = None),

prelude/rust/rust_binary.bzl

@@ -359,7 +359,7 @@ def rust_test_impl(ctx: AnalysisContext) -> list[[DefaultInfo.type, RunInfo.type
     )
 
     # Setup a RE executor based on the `remote_execution` param.
-    re_executor = get_re_executor_from_props(ctx.attrs.remote_execution)
+    re_executor = get_re_executor_from_props(ctx)
 
     return inject_test_run_info(
         ctx,

prelude/sh_test.bzl

@@ -42,7 +42,7 @@ def sh_test_impl(ctx: AnalysisContext) -> list[Provider]:
     command = [args] + ctx.attrs.args
 
     # Setup a RE executor based on the `remote_execution` param.
-    re_executor = get_re_executor_from_props(ctx.attrs.remote_execution)
+    re_executor = get_re_executor_from_props(ctx)
 
     # We implicitly make the target run from the project root if remote
     # excution options were specified