From 5e88501cb33f44d3b3c9fb77f94832699b05b4c7 Mon Sep 17 00:00:00 2001
From: Richard Levasseur <rlevasseur@google.com>
Date: Sat, 24 May 2025 12:08:39 -0700
Subject: [PATCH] tests: refacto py_reconfig rules so less boilerplate is
 needed to add attrs

---
 tests/support/sh_py_run_test.bzl | 82 ++++++++++++++++++++++++++++++--
 1 file changed, 78 insertions(+), 4 deletions(-)

diff --git a/tests/support/sh_py_run_test.bzl b/tests/support/sh_py_run_test.bzl
index 1a61de9bd3..04a2883fde 100644
--- a/tests/support/sh_py_run_test.bzl
+++ b/tests/support/sh_py_run_test.bzl
@@ -13,14 +13,88 @@
 # limitations under the License.
 """Run a py_binary with altered config settings in an sh_test.
 
-This facilitates verify running binaries with different outer environmental
-settings and verifying their output without the overhead of a bazel-in-bazel
-integration test.
+This facilitates verify running binaries with different configuration settings
+without the overhead of a bazel-in-bazel integration test.
 """
 
 load("@rules_shell//shell:sh_test.bzl", "sh_test")
+load("//python/private:attr_builders.bzl", "attrb")  # buildifier: disable=bzl-visibility
+load("//python/private:py_binary_macro.bzl", "py_binary_macro")  # buildifier: disable=bzl-visibility
+load("//python/private:py_binary_rule.bzl", "create_py_binary_rule_builder")  # buildifier: disable=bzl-visibility
+load("//python/private:py_test_macro.bzl", "py_test_macro")  # buildifier: disable=bzl-visibility
+load("//python/private:py_test_rule.bzl", "create_py_test_rule_builder")  # buildifier: disable=bzl-visibility
 load("//python/private:toolchain_types.bzl", "TARGET_TOOLCHAIN_TYPE")  # buildifier: disable=bzl-visibility
-load(":py_reconfig.bzl", "py_reconfig_binary")
+load("//tests/support:support.bzl", "VISIBLE_FOR_TESTING")
+
+def _perform_transition_impl(input_settings, attr, base_impl):
+    settings = {k: input_settings[k] for k in _RECONFIG_INHERITED_OUTPUTS if k in input_settings}
+    settings.update(base_impl(input_settings, attr))
+
+    settings[VISIBLE_FOR_TESTING] = True
+    settings["//command_line_option:build_python_zip"] = attr.build_python_zip
+
+    for attr_name, setting_label in _RECONFIG_ATTR_SETTING_MAP.items():
+        if getattr(attr, attr_name):
+            settings[setting_label] = getattr(attr, attr_name)
+    return settings
+
+# Attributes that, if non-falsey (`if attr.<name>`), will copy their
+# value into the output settings
+_RECONFIG_ATTR_SETTING_MAP = {
+    "bootstrap_impl": "//python/config_settings:bootstrap_impl",
+    "extra_toolchains": "//command_line_option:extra_toolchains",
+    "python_src": "//python/bin:python_src",
+    "venvs_site_packages": "//python/config_settings:venvs_site_packages",
+    "venvs_use_declare_symlink": "//python/config_settings:venvs_use_declare_symlink",
+}
+
+_RECONFIG_INPUTS = _RECONFIG_ATTR_SETTING_MAP.values()
+_RECONFIG_OUTPUTS = _RECONFIG_INPUTS + [
+    "//command_line_option:build_python_zip",
+    VISIBLE_FOR_TESTING,
+]
+_RECONFIG_INHERITED_OUTPUTS = [v for v in _RECONFIG_OUTPUTS if v in _RECONFIG_INPUTS]
+
+_RECONFIG_ATTRS = {
+    "bootstrap_impl": attrb.String(),
+    "build_python_zip": attrb.String(default = "auto"),
+    "extra_toolchains": attrb.StringList(
+        doc = """
+Value for the --extra_toolchains flag.
+
+NOTE: You'll likely have to also specify //tests/support/cc_toolchains:all (or some CC toolchain)
+to make the RBE presubmits happy, which disable auto-detection of a CC
+toolchain.
+""",
+    ),
+    "python_src": attrb.Label(),
+    "venvs_site_packages": attrb.String(),
+    "venvs_use_declare_symlink": attrb.String(),
+}
+
+def _create_reconfig_rule(builder):
+    builder.attrs.update(_RECONFIG_ATTRS)
+
+    base_cfg_impl = builder.cfg.implementation()
+    builder.cfg.set_implementation(lambda *args: _perform_transition_impl(base_impl = base_cfg_impl, *args))
+    builder.cfg.update_inputs(_RECONFIG_INPUTS)
+    builder.cfg.update_outputs(_RECONFIG_OUTPUTS)
+    return builder.build()
+
+_py_reconfig_binary = _create_reconfig_rule(create_py_binary_rule_builder())
+
+_py_reconfig_test = _create_reconfig_rule(create_py_test_rule_builder())
+
+def py_reconfig_test(**kwargs):
+    """Create a py_test with customized build settings for testing.
+
+    Args:
+        **kwargs: kwargs to pass along to _py_reconfig_test.
+    """
+    py_test_macro(_py_reconfig_test, **kwargs)
+
+def py_reconfig_binary(**kwargs):
+    py_binary_macro(_py_reconfig_binary, **kwargs)
 
 def sh_py_run_test(*, name, sh_src, py_src, **kwargs):
     """Run a py_binary within a sh_test.