From db62c7753a6bb941934dc55734a3aa4c3077c3d2 Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Mon, 14 Oct 2024 15:22:09 +0200
Subject: [PATCH] `platform_transition_test`: Use transitioned target platform

The target platform of a test rule matters for the resolution of the execution platform of the test action, either via test toolchains or `--use_target_platform_for_tests`. With an outgoing transition, tests would run on the wrong platform in these cases, so use an incoming transition for the test rule.
---
 lib/tests/transitions/BUILD.bazel | 15 ++++-------
 lib/transitions.bzl               | 41 ++++++++++++++++++++-----------
 2 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/lib/tests/transitions/BUILD.bazel b/lib/tests/transitions/BUILD.bazel
index 684d0c50d..48abfc89c 100644
--- a/lib/tests/transitions/BUILD.bazel
+++ b/lib/tests/transitions/BUILD.bazel
@@ -1,5 +1,6 @@
 load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library", "go_test")
+load("@platforms//host:constraints.bzl", "HOST_CONSTRAINTS")
 load("//lib:diff_test.bzl", "diff_test")
 load(
     "//lib:transitions.bzl",
@@ -135,22 +136,16 @@ platform_transition_binary(
 platform_transition_test(
     name = "transitioned_go_test_x86_64",
     binary = ":test_transition_test",
-    # only run this test on x86_64 platforms
-    target_compatible_with = [
-        "@platforms//cpu:x86_64",
-        "@platforms//os:linux",
-    ],
+    # only run this test on an x86_64 host
+    target_compatible_with = [] if sorted(HOST_CONSTRAINTS) == ["@platforms//cpu:x86_64", "@platforms//os:linux"] else ["@platforms//:incompatible"],
     target_platform = "x86_64_linux",
 )
 
 platform_transition_test(
     name = "transitioned_go_test_arm64",
     binary = ":test_transition_test",
-    # only run this test on arm64 platforms
-    target_compatible_with = [
-        "@platforms//cpu:arm64",
-        "@platforms//os:linux",
-    ],
+    # only run this test on an x86_64 host
+    target_compatible_with = [] if sorted(HOST_CONSTRAINTS) == ["@platforms//cpu:aarch64", "@platforms//os:linux"] else ["@platforms//:incompatible"],
     target_platform = "arm64_linux",
 )
 
diff --git a/lib/transitions.bzl b/lib/transitions.bzl
index 9fda22c58..59cbd48f6 100644
--- a/lib/transitions.bzl
+++ b/lib/transitions.bzl
@@ -52,7 +52,12 @@ def _platform_transition_binary_impl(ctx):
     # forwarding DefaultInfo.
 
     result = []
-    binary = ctx.attr.binary[0]
+
+    # ctx.attr.binary is a singleton list if this rule uses an outgoing transition.
+    if type(ctx.attr.binary) == type([]):
+        binary = ctx.attr.binary[0]
+    else:
+        binary = ctx.attr.binary
 
     default_info = binary[DefaultInfo]
     files = default_info.files
@@ -87,28 +92,36 @@ def _platform_transition_binary_impl(ctx):
 
     return result
 
-_platform_transition_attrs = {
-    "basename": attr.string(),
-    "binary": attr.label(allow_files = True, cfg = _transition_platform),
-    "target_platform": attr.label(
-        doc = "The target platform to transition the binary.",
-        mandatory = True,
-    ),
-    "_allowlist_function_transition": attr.label(
-        default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
-    ),
-}
+def _get_platform_transition_attrs(binary_cfg):
+    return {
+        "basename": attr.string(),
+        "binary": attr.label(allow_files = True, cfg = binary_cfg),
+        "target_platform": attr.label(
+            doc = "The target platform to transition the binary.",
+            mandatory = True,
+        ),
+        "_allowlist_function_transition": attr.label(
+            default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
+        ),
+    }
 
 platform_transition_binary = rule(
     implementation = _platform_transition_binary_impl,
-    attrs = _platform_transition_attrs,
+    # Use an outgoing transition since the target platform of the
+    # platform_transition_binary doesn't matter and it results in a more
+    # intuitive output path (matching an untransitioned binary).
+    attrs = _get_platform_transition_attrs(binary_cfg = _transition_platform),
     executable = True,
     doc = "Transitions the binary to use the provided platform.",
 )
 
 platform_transition_test = rule(
     implementation = _platform_transition_binary_impl,
-    attrs = _platform_transition_attrs,
+    attrs = _get_platform_transition_attrs(binary_cfg = "target"),
+    # Use an incoming transition since the target platform of the
+    # platform_transition_test does matter for the exec platform resolution of
+    # the test action.
+    cfg = _transition_platform,
     test = True,
     doc = "Transitions the test to use the provided platform.",
 )