tests: make toolchain tests run in the same build instead of bazel-in-bazel integration tests (bazel-contrib#2095)

This changes the //tests/toolchains tests to run in the same Bazel build
instance instead of
being bazel-in-bazel integration tests. This is done by using a
transition to set
the python version to match the desired toolchain.

This makes running the tests **much** cheaper -- there were 37 sub-bazel
processes being
started (one for each python version), which was very expensive.
Debugging is also much easier because they aren't part of a
bazel-in-bazel invocation.

Author: rickeylev

MODULE.bazel

@@ -76,6 +76,15 @@ bazel_dep(name = "rules_testing", version = "0.6.0", dev_dependency = True)
 bazel_dep(name = "rules_go", version = "0.41.0", dev_dependency = True, repo_name = "io_bazel_rules_go")
 bazel_dep(name = "gazelle", version = "0.33.0", dev_dependency = True, repo_name = "bazel_gazelle")
 
+dev_python = use_extension(
+    "//python/extensions:python.bzl",
+    "python",
+    dev_dependency = True,
+)
+dev_python.rules_python_private_testing(
+    register_all_versions = True,
+)
+
 dev_pip = use_extension(
     "//python/private/pypi:pip.bzl",
     "pip_internal",

WORKSPACE

@@ -42,12 +42,13 @@ load("//:internal_setup.bzl", "rules_python_internal_setup")
 rules_python_internal_setup()
 
 load("//python:repositories.bzl", "python_register_multi_toolchains")
-load("//python:versions.bzl", "MINOR_MAPPING")
+load("//python:versions.bzl", "MINOR_MAPPING", "TOOL_VERSIONS")
 
 python_register_multi_toolchains(
     name = "python",
     default_version = MINOR_MAPPING.values()[-2],
-    python_versions = MINOR_MAPPING.values(),
+    # Integration tests verify each version, so register all of them.
+    python_versions = TOOL_VERSIONS.keys(),
 )
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive", "http_file")

python/private/python.bzl

@@ -16,6 +16,7 @@
 
 load("@bazel_features//:features.bzl", "bazel_features")
 load("//python:repositories.bzl", "python_register_toolchains")
+load("//python:versions.bzl", "TOOL_VERSIONS")
 load(":pythons_hub.bzl", "hub_repo")
 load(":text_util.bzl", "render")
 load(":toolchains_repo.bzl", "multi_toolchain_aliases")
@@ -78,7 +79,9 @@ def _python_impl(module_ctx):
     for mod in module_ctx.modules:
         module_toolchain_versions = []
 
-        for toolchain_attr in mod.tags.toolchain:
+        toolchain_attr_structs = _create_toolchain_attr_structs(mod)
+
+        for toolchain_attr in toolchain_attr_structs:
             toolchain_version = toolchain_attr.python_version
             toolchain_name = "python_" + toolchain_version.replace(".", "_")
 
@@ -95,9 +98,7 @@ def _python_impl(module_ctx):
                 # * rules_python needs to set a soft default in case the root module doesn't,
                 #   e.g. if the root module doesn't use Python itself.
                 # * The root module is allowed to override the rules_python default.
-
-                # A single toolchain is treated as the default because it's unambiguous.
-                is_default = toolchain_attr.is_default or len(mod.tags.toolchain) == 1
+                is_default = toolchain_attr.is_default
 
                 # Also only the root module should be able to decide ignore_root_user_error.
                 # Modules being depended upon don't know the final environment, so they aren't
@@ -251,6 +252,43 @@ def _fail_multiple_default_toolchains(first, second):
         second = second,
     ))
 
+def _create_toolchain_attr_structs(mod):
+    arg_structs = []
+    seen_versions = {}
+    for tag in mod.tags.toolchain:
+        arg_structs.append(_create_toolchain_attrs_struct(tag = tag, toolchain_tag_count = len(mod.tags.toolchain)))
+        seen_versions[tag.python_version] = True
+
+    if mod.is_root:
+        register_all = False
+        for tag in mod.tags.rules_python_private_testing:
+            if tag.register_all_versions:
+                register_all = True
+                break
+        if register_all:
+            arg_structs.extend([
+                _create_toolchain_attrs_struct(python_version = v)
+                for v in TOOL_VERSIONS.keys()
+                if v not in seen_versions
+            ])
+    return arg_structs
+
+def _create_toolchain_attrs_struct(*, tag = None, python_version = None, toolchain_tag_count = None):
+    if tag and python_version:
+        fail("Only one of tag and python version can be specified")
+    if tag:
+        # A single toolchain is treated as the default because it's unambiguous.
+        is_default = tag.is_default or toolchain_tag_count == 1
+    else:
+        is_default = False
+
+    return struct(
+        is_default = is_default,
+        python_version = python_version if python_version else tag.python_version,
+        configure_coverage_tool = getattr(tag, "configure_coverage_tool", False),
+        ignore_root_user_error = getattr(tag, "ignore_root_user_error", False),
+    )
+
 def _get_bazel_version_specific_kwargs():
     kwargs = {}
 
@@ -264,6 +302,11 @@ python = module_extension(
 """,
     implementation = _python_impl,
     tag_classes = {
+        "rules_python_private_testing": tag_class(
+            attrs = {
+                "register_all_versions": attr.bool(default = False),
+            },
+        ),
         "toolchain": tag_class(
             doc = """Tag class used to register Python toolchains.
 Use this tag class to register one or more Python toolchains. This class

tests/support/sh_py_run_test.bzl

@@ -20,28 +20,33 @@ without the overhead of a bazel-in-bazel integration test.
 load("//python:py_binary.bzl", "py_binary")
 load("//python:py_test.bzl", "py_test")
 load("//python/private:toolchain_types.bzl", "TARGET_TOOLCHAIN_TYPE")  # buildifier: disable=bzl-visibility
+load("//tests/support:support.bzl", "VISIBLE_FOR_TESTING")
 
 def _perform_transition_impl(input_settings, attr):
     settings = dict(input_settings)
+    settings[VISIBLE_FOR_TESTING] = True
     settings["//command_line_option:build_python_zip"] = attr.build_python_zip
     if attr.bootstrap_impl:
         settings["//python/config_settings:bootstrap_impl"] = attr.bootstrap_impl
     if attr.extra_toolchains:
         settings["//command_line_option:extra_toolchains"] = attr.extra_toolchains
-    else:
-        settings["//command_line_option:extra_toolchains"] = input_settings["//command_line_option:extra_toolchains"]
+    if attr.python_version:
+        settings["//python/config_settings:python_version"] = attr.python_version
     return settings
 
 _perform_transition = transition(
     implementation = _perform_transition_impl,
     inputs = [
         "//python/config_settings:bootstrap_impl",
         "//command_line_option:extra_toolchains",
+        "//python/config_settings:python_version",
     ],
     outputs = [
         "//command_line_option:build_python_zip",
         "//command_line_option:extra_toolchains",
         "//python/config_settings:bootstrap_impl",
+        "//python/config_settings:python_version",
+        VISIBLE_FOR_TESTING,
     ],
 )
 
@@ -99,6 +104,7 @@ to make the RBE presubmits happy, which disable auto-detection of a CC
 toolchain.
 """,
         ),
+        "python_version": attr.string(),
         "target": attr.label(executable = True, cfg = "target"),
         "_allowlist_function_transition": attr.label(
             default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
@@ -125,7 +131,10 @@ def py_reconfig_test(*, name, **kwargs):
     reconfig_kwargs = {}
     reconfig_kwargs["bootstrap_impl"] = kwargs.pop("bootstrap_impl", None)
     reconfig_kwargs["extra_toolchains"] = kwargs.pop("extra_toolchains", None)
+    reconfig_kwargs["python_version"] = kwargs.pop("python_version", None)
     reconfig_kwargs["env"] = kwargs.get("env")
+    reconfig_kwargs["target_compatible_with"] = kwargs.get("target_compatible_with")
+
     inner_name = "_{}_inner" + name
     _py_reconfig_test(
         name = name,
@@ -178,6 +187,7 @@ def _current_build_settings_impl(ctx):
                 "short_path": runtime.interpreter.short_path if runtime.interpreter else None,
             },
             "interpreter_path": runtime.interpreter_path,
+            "toolchain_label": str(getattr(toolchain, "toolchain_label", None)),
         }),
     )
     return [DefaultInfo(

tests/toolchains/BUILD.bazel

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load(":defs.bzl", "acceptance_tests")
-load(":versions_test.bzl", "versions_test_suite")
+load(":defs.bzl", "define_toolchain_tests")
 
-versions_test_suite(name = "versions_test")
-
-acceptance_tests()
+define_toolchain_tests(
+    name = "toolchain_tests",
+)