internal: support repo prefixes and config settings in alias rendering (#1756)

This does not change any logic/features in the bzlmod or legacy code,
but
just changes the interfaces and how the parameters are passed. The final
result should be the same.

Summary:
- introduce a `whl_alias` struct to make code more object oriented.
- make the interface of the function as small as possible.
- unify the bzlmod and legacy code paths and simplify tests.
- allow to specify arbitrary repos and config settings
  when generating aliases, which is useful in multi-platform hub
  generation.

---------

Co-authored-by: Richard Levasseur <richardlev@gmail.com>
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
index 110ade1..7b8160e 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -26,7 +26,7 @@
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:parse_whl_name.bzl", "parse_whl_name")
 load("//python/private:patch_whl.bzl", "patch_whl")
-load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")
+load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")
 load("//python/private:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
 load("//python/private:toolchains_repo.bzl", "get_host_os_arch")
 load("//python/private:whl_target_platforms.bzl", "whl_target_platforms")
@@ -374,7 +374,13 @@
         config["experimental_target_platforms"] = rctx.attr.experimental_target_platforms
 
     macro_tmpl = "@%s//{}:{}" % rctx.attr.name
-    aliases = render_pkg_aliases(repo_name = rctx.attr.name, bzl_packages = bzl_packages)
+
+    aliases = render_pkg_aliases(
+        aliases = {
+            pkg: [whl_alias(repo = rctx.attr.name + "_" + pkg)]
+            for pkg in bzl_packages or []
+        },
+    )
     for path, contents in aliases.items():
         rctx.file(path, contents)
 
diff --git a/python/private/bzlmod/pip.bzl b/python/private/bzlmod/pip.bzl
index b4dbf2f..a017089 100644
--- a/python/private/bzlmod/pip.bzl
+++ b/python/private/bzlmod/pip.bzl
@@ -27,6 +27,7 @@
 load("//python/pip_install:requirements_parser.bzl", parse_requirements = "parse")
 load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:parse_whl_name.bzl", "parse_whl_name")
+load("//python/private:render_pkg_aliases.bzl", "whl_alias")
 load("//python/private:version_label.bzl", "version_label")
 load(":pip_repository.bzl", "pip_repository")
 
@@ -179,8 +180,9 @@
         group_name = whl_group_mapping.get(whl_name)
         group_deps = requirement_cycles.get(group_name, [])
 
+        repo_name = "{}_{}".format(pip_name, whl_name)
         whl_library(
-            name = "%s_%s" % (pip_name, whl_name),
+            name = repo_name,
             requirement = requirement_line,
             repo = pip_name,
             repo_prefix = pip_name + "_",
@@ -205,10 +207,15 @@
             group_deps = group_deps,
         )
 
-        if whl_name not in whl_map[hub_name]:
-            whl_map[hub_name][whl_name] = {}
-
-        whl_map[hub_name][whl_name][_major_minor_version(pip_attr.python_version)] = pip_name + "_"
+        major_minor = _major_minor_version(pip_attr.python_version)
+        whl_map[hub_name].setdefault(whl_name, []).append(
+            whl_alias(
+                repo = repo_name,
+                version = major_minor,
+                # Call Label() to canonicalize because its used in a different context
+                config_setting = Label("//python/config_settings:is_python_" + major_minor),
+            ),
+        )
 
 def _pip_impl(module_ctx):
     """Implementation of a class tag that creates the pip hub and corresponding pip spoke whl repositories.
@@ -358,7 +365,10 @@
         pip_repository(
             name = hub_name,
             repo_name = hub_name,
-            whl_map = whl_map,
+            whl_map = {
+                key: json.encode(value)
+                for key, value in whl_map.items()
+            },
             default_version = _major_minor_version(DEFAULT_PYTHON_VERSION),
         )
 
diff --git a/python/private/bzlmod/pip_repository.bzl b/python/private/bzlmod/pip_repository.bzl
index 8ea5ee7..d96131d 100644
--- a/python/private/bzlmod/pip_repository.bzl
+++ b/python/private/bzlmod/pip_repository.bzl
@@ -14,7 +14,7 @@
 
 ""
 
-load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")
+load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")
 load("//python/private:text_util.bzl", "render")
 
 _BUILD_FILE_CONTENTS = """\
@@ -27,10 +27,11 @@
 def _pip_repository_impl(rctx):
     bzl_packages = rctx.attr.whl_map.keys()
     aliases = render_pkg_aliases(
-        repo_name = rctx.attr.repo_name,
-        rules_python = rctx.attr._template.workspace_name,
+        aliases = {
+            key: [whl_alias(**v) for v in json.decode(values)]
+            for key, values in rctx.attr.whl_map.items()
+        },
         default_version = rctx.attr.default_version,
-        whl_map = rctx.attr.whl_map,
     )
     for path, contents in aliases.items():
         rctx.file(path, contents)
@@ -71,9 +72,12 @@
         mandatory = True,
         doc = "The apparent name of the repo. This is needed because in bzlmod, the name attribute becomes the canonical name.",
     ),
-    "whl_map": attr.string_list_dict(
+    "whl_map": attr.string_dict(
         mandatory = True,
-        doc = "The wheel map where values are python versions",
+        doc = """\
+The wheel map where values are json.encoded strings of the whl_map constructed
+in the pip.parse tag class.
+""",
     ),
     "_template": attr.label(
         default = ":requirements.bzl.tmpl",
diff --git a/python/private/render_pkg_aliases.bzl b/python/private/render_pkg_aliases.bzl
index 02ba75b..e38f133 100644
--- a/python/private/render_pkg_aliases.bzl
+++ b/python/private/render_pkg_aliases.bzl
@@ -16,9 +16,8 @@
 
 This is used in bzlmod and non-bzlmod setups."""
 
-load("//python/private:normalize_name.bzl", "normalize_name")
+load(":normalize_name.bzl", "normalize_name")
 load(":text_util.bzl", "render")
-load(":version_label.bzl", "version_label")
 
 NO_MATCH_ERROR_MESSAGE_TEMPLATE = """\
 No matching wheel for current configuration's Python version.
@@ -42,56 +41,31 @@
 def _render_whl_library_alias(
         *,
         name,
-        repo_name,
-        dep,
-        target,
         default_version,
-        versions,
-        rules_python):
-    """Render an alias for common targets
-
-    If the versions is passed, then the `rules_python` must be passed as well and
-    an alias with a select statement based on the python version is going to be
-    generated.
-    """
-    if versions == None:
+        aliases):
+    """Render an alias for common targets."""
+    if len(aliases) == 1 and not aliases[0].version:
+        alias = aliases[0]
         return render.alias(
             name = name,
-            actual = repr("@{repo_name}_{dep}//:{target}".format(
-                repo_name = repo_name,
-                dep = dep,
-                target = target,
-            )),
+            actual = repr("@{repo}//:{name}".format(repo = alias.repo, name = name)),
         )
 
     # Create the alias repositories which contains different select
     # statements  These select statements point to the different pip
     # whls that are based on a specific version of Python.
     selects = {}
-    for full_version in versions:
-        condition = "@@{rules_python}//python/config_settings:is_python_{full_python_version}".format(
-            rules_python = rules_python,
-            full_python_version = full_version,
-        )
-        actual = "@{repo_name}_{version}_{dep}//:{target}".format(
-            repo_name = repo_name,
-            version = version_label(full_version),
-            dep = dep,
-            target = target,
-        )
-        selects[condition] = actual
+    no_match_error = "_NO_MATCH_ERROR"
+    default = None
+    for alias in sorted(aliases, key = lambda x: x.version):
+        actual = "@{repo}//:{name}".format(repo = alias.repo, name = name)
+        selects[alias.config_setting] = actual
+        if alias.version == default_version:
+            default = actual
+            no_match_error = None
 
-    if default_version:
-        no_match_error = None
-        default_actual = "@{repo_name}_{version}_{dep}//:{target}".format(
-            repo_name = repo_name,
-            version = version_label(default_version),
-            dep = dep,
-            target = target,
-        )
-        selects["//conditions:default"] = default_actual
-    else:
-        no_match_error = "_NO_MATCH_ERROR"
+    if default:
+        selects["//conditions:default"] = default
 
     return render.alias(
         name = name,
@@ -101,22 +75,21 @@
         ),
     )
 
-def _render_common_aliases(repo_name, name, versions = None, default_version = None, rules_python = None):
+def _render_common_aliases(*, name, aliases, default_version = None):
     lines = [
         """package(default_visibility = ["//visibility:public"])""",
     ]
 
-    if versions:
-        versions = sorted(versions)
+    versions = None
+    if aliases:
+        versions = sorted([v.version for v in aliases if v.version])
 
-    if not versions:
-        pass
-    elif default_version in versions:
+    if not versions or default_version in versions:
         pass
     else:
         error_msg = NO_MATCH_ERROR_MESSAGE_TEMPLATE.format(
             supported_versions = ", ".join(versions),
-            rules_python = rules_python,
+            rules_python = "rules_python",
         )
 
         lines.append("_NO_MATCH_ERROR = \"\"\"\\\n{error_msg}\"\"\"".format(
@@ -137,12 +110,8 @@
         [
             _render_whl_library_alias(
                 name = target,
-                repo_name = repo_name,
-                dep = name,
-                target = target,
-                versions = versions,
                 default_version = default_version,
-                rules_python = rules_python,
+                aliases = aliases,
             )
             for target in ["pkg", "whl", "data", "dist_info"]
         ],
@@ -150,7 +119,7 @@
 
     return "\n\n".join(lines)
 
-def render_pkg_aliases(*, repo_name, bzl_packages = None, whl_map = None, rules_python = None, default_version = None):
+def render_pkg_aliases(*, aliases, default_version = None):
     """Create alias declarations for each PyPI package.
 
     The aliases should be appended to the pip_repository BUILD.bazel file. These aliases
@@ -158,36 +127,55 @@
     when using bzlmod.
 
     Args:
-        repo_name: the repository name of the hub repository that is visible to the users that is
-            also used as the prefix for the spoke repo names (e.g. "pip", "pypi").
-        bzl_packages: the list of packages to setup, if not specified, whl_map.keys() will be used instead.
-        whl_map: the whl_map for generating Python version aware aliases.
+        aliases: dict, the keys are normalized distribution names and values are the
+            whl_alias instances.
         default_version: the default version to be used for the aliases.
-        rules_python: the name of the rules_python workspace.
 
     Returns:
         A dict of file paths and their contents.
     """
-    if not bzl_packages and whl_map:
-        bzl_packages = list(whl_map.keys())
-
     contents = {}
-    if not bzl_packages:
+    if not aliases:
         return contents
+    elif type(aliases) != type({}):
+        fail("The aliases need to be provided as a dict, got: {}".format(type(aliases)))
 
-    for name in bzl_packages:
-        versions = None
-        if whl_map != None:
-            versions = whl_map[name]
-        name = normalize_name(name)
-
-        filename = "{}/BUILD.bazel".format(name)
-        contents[filename] = _render_common_aliases(
-            repo_name = repo_name,
-            name = name,
-            versions = versions,
-            rules_python = rules_python,
+    return {
+        "{}/BUILD.bazel".format(normalize_name(name)): _render_common_aliases(
+            name = normalize_name(name),
+            aliases = pkg_aliases,
             default_version = default_version,
         ).strip()
+        for name, pkg_aliases in aliases.items()
+    }
 
-    return contents
+def whl_alias(*, repo, version = None, config_setting = None):
+    """The bzl_packages value used by by the render_pkg_aliases function.
+
+    This contains the minimum amount of information required to generate correct
+    aliases in a hub repository.
+
+    Args:
+        repo: str, the repo of where to find the things to be aliased.
+        version: optional(str), the version of the python toolchain that this
+            whl alias is for. If not set, then non-version aware aliases will be
+            constructed. This is mainly used for better error messages when there
+            is no match found during a select.
+        config_setting: optional(Label or str), the config setting that we should use. Defaults
+            to "@rules_python//python/config_settings:is_python_{version}".
+
+    Returns:
+        a struct with the validated and parsed values.
+    """
+    if not repo:
+        fail("'repo' must be specified")
+
+    if version:
+        config_setting = config_setting or Label("//python/config_settings:is_python_" + version)
+        config_setting = str(config_setting)
+
+    return struct(
+        repo = repo,
+        version = version,
+        config_setting = config_setting,
+    )
diff --git a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
index 513c278..c61e5ef 100644
--- a/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
+++ b/tests/pip_hub_repository/render_pkg_aliases/render_pkg_aliases_test.bzl
@@ -15,14 +15,37 @@
 """render_pkg_aliases tests"""
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
-load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")  # buildifier: disable=bzl-visibility
+load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
+load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases", "whl_alias")  # buildifier: disable=bzl-visibility
+
+def _normalize_label_strings(want):
+    """normalize expected strings.
+
+    This function ensures that the desired `render_pkg_aliases` outputs are
+    normalized from `bzlmod` to `WORKSPACE` values so that we don't have to
+    have to sets of expected strings. The main difference is that under
+    `bzlmod` the `str(Label("//my_label"))` results in `"@@//my_label"` whereas
+    under `non-bzlmod` we have `"@//my_label"`. This function does
+    `string.replace("@@", "@")` to normalize the strings.
+
+    NOTE, in tests, we should only use keep `@@` usage in expectation values
+    for the test cases where the whl_alias has the `config_setting` constructed
+    from a `Label` instance.
+    """
+    if "@@" not in want:
+        fail("The expected string does not have '@@' labels, consider not using the function")
+
+    if BZLMOD_ENABLED:
+        # our expectations are already with double @
+        return want
+
+    return want.replace("@@", "@")
 
 _tests = []
 
 def _test_empty(env):
     actual = render_pkg_aliases(
-        bzl_packages = None,
-        repo_name = "pypi",
+        aliases = None,
     )
 
     want = {}
@@ -33,12 +56,15 @@
 
 def _test_legacy_aliases(env):
     actual = render_pkg_aliases(
-        bzl_packages = ["foo"],
-        repo_name = "pypi",
+        aliases = {
+            "foo": [
+                whl_alias(repo = "pypi_foo"),
+            ],
+        },
     )
 
-    want = {
-        "foo/BUILD.bazel": """\
+    want_key = "foo/BUILD.bazel"
+    want_content = """\
 package(default_visibility = ["//visibility:public"])
 
 alias(
@@ -64,37 +90,24 @@
 alias(
     name = "dist_info",
     actual = "@pypi_foo//:dist_info",
-)""",
-    }
+)"""
 
-    env.expect.that_dict(actual).contains_exactly(want)
+    env.expect.that_dict(actual).contains_exactly({want_key: want_content})
 
 _tests.append(_test_legacy_aliases)
 
-def _test_all_legacy_aliases_are_created(env):
-    actual = render_pkg_aliases(
-        bzl_packages = ["foo", "bar"],
-        repo_name = "pypi",
-    )
-
-    want_files = ["bar/BUILD.bazel", "foo/BUILD.bazel"]
-
-    env.expect.that_dict(actual).keys().contains_exactly(want_files)
-
-_tests.append(_test_all_legacy_aliases_are_created)
-
 def _test_bzlmod_aliases(env):
     actual = render_pkg_aliases(
         default_version = "3.2",
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar-baz": ["3.2"],
+        aliases = {
+            "bar-baz": [
+                whl_alias(version = "3.2", repo = "pypi_32_bar_baz", config_setting = "//:my_config_setting"),
+            ],
         },
     )
 
-    want = {
-        "bar_baz/BUILD.bazel": """\
+    want_key = "bar_baz/BUILD.bazel"
+    want_content = """\
 package(default_visibility = ["//visibility:public"])
 
 alias(
@@ -106,7 +119,7 @@
     name = "pkg",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "//:my_config_setting": "@pypi_32_bar_baz//:pkg",
             "//conditions:default": "@pypi_32_bar_baz//:pkg",
         },
     ),
@@ -116,7 +129,7 @@
     name = "whl",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "//:my_config_setting": "@pypi_32_bar_baz//:whl",
             "//conditions:default": "@pypi_32_bar_baz//:whl",
         },
     ),
@@ -126,7 +139,7 @@
     name = "data",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "//:my_config_setting": "@pypi_32_bar_baz//:data",
             "//conditions:default": "@pypi_32_bar_baz//:data",
         },
     ),
@@ -136,24 +149,30 @@
     name = "dist_info",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "//:my_config_setting": "@pypi_32_bar_baz//:dist_info",
             "//conditions:default": "@pypi_32_bar_baz//:dist_info",
         },
     ),
-)""",
-    }
+)"""
 
-    env.expect.that_dict(actual).contains_exactly(want)
+    env.expect.that_collection(actual.keys()).contains_exactly([want_key])
+    env.expect.that_str(actual[want_key]).equals(want_content)
 
 _tests.append(_test_bzlmod_aliases)
 
 def _test_bzlmod_aliases_with_no_default_version(env):
     actual = render_pkg_aliases(
         default_version = None,
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar-baz": ["3.2", "3.1"],
+        aliases = {
+            "bar-baz": [
+                whl_alias(
+                    version = "3.2",
+                    repo = "pypi_32_bar_baz",
+                    # pass the label to ensure that it gets converted to string
+                    config_setting = Label("//python/config_settings:is_python_3.2"),
+                ),
+                whl_alias(version = "3.1", repo = "pypi_31_bar_baz"),
+            ],
         },
     )
 
@@ -189,8 +208,8 @@
     name = "pkg",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -200,8 +219,8 @@
     name = "whl",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -211,8 +230,8 @@
     name = "data",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -222,15 +241,15 @@
     name = "dist_info",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(want_content)
+    env.expect.that_str(actual[want_key]).equals(_normalize_label_strings(want_content))
 
 _tests.append(_test_bzlmod_aliases_with_no_default_version)
 
@@ -244,10 +263,11 @@
         # non-root module, then we will have a no-match-error because the default_version
         # is not in the list of the versions in the whl_map.
         default_version = "3.3",
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar-baz": ["3.2", "3.1"],
+        aliases = {
+            "bar-baz": [
+                whl_alias(version = "3.2", repo = "pypi_32_bar_baz"),
+                whl_alias(version = "3.1", repo = "pypi_31_bar_baz"),
+            ],
         },
     )
 
@@ -283,8 +303,8 @@
     name = "pkg",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:pkg",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:pkg",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -294,8 +314,8 @@
     name = "whl",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:whl",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:whl",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -305,8 +325,8 @@
     name = "data",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:data",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:data",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
@@ -316,26 +336,30 @@
     name = "dist_info",
     actual = select(
         {
-            "@@rules_python//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
-            "@@rules_python//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.1": "@pypi_31_bar_baz//:dist_info",
+            "@@//python/config_settings:is_python_3.2": "@pypi_32_bar_baz//:dist_info",
         },
         no_match_error = _NO_MATCH_ERROR,
     ),
 )"""
 
     env.expect.that_collection(actual.keys()).contains_exactly([want_key])
-    env.expect.that_str(actual[want_key]).equals(want_content)
+    env.expect.that_str(actual[want_key]).equals(_normalize_label_strings(want_content))
 
 _tests.append(_test_bzlmod_aliases_for_non_root_modules)
 
-def _test_bzlmod_aliases_are_created_for_all_wheels(env):
+def _test_aliases_are_created_for_all_wheels(env):
     actual = render_pkg_aliases(
         default_version = "3.2",
-        repo_name = "pypi",
-        rules_python = "rules_python",
-        whl_map = {
-            "bar": ["3.1", "3.2"],
-            "foo": ["3.1", "3.2"],
+        aliases = {
+            "bar": [
+                whl_alias(version = "3.1", repo = "pypi_31_bar"),
+                whl_alias(version = "3.2", repo = "pypi_32_bar"),
+            ],
+            "foo": [
+                whl_alias(version = "3.1", repo = "pypi_32_foo"),
+                whl_alias(version = "3.2", repo = "pypi_31_foo"),
+            ],
         },
     )
 
@@ -346,7 +370,7 @@
 
     env.expect.that_dict(actual).keys().contains_exactly(want_files)
 
-_tests.append(_test_bzlmod_aliases_are_created_for_all_wheels)
+_tests.append(_test_aliases_are_created_for_all_wheels)
 
 def render_pkg_aliases_test_suite(name):
     """Create the test suite.