test(bzlmod): Make some tests bzlmod compatible with Bazel@HEAD (#1482)

A few tests weren't compatible with bzlmod, so would fail when it was
enabled. The various causes and fixes are:

* Under bzlmod, `runfiles.CurrentRepository()` returns the empty string
for the main repository. To fix, an environment variable is used to tell
the test whether bzlmod is enabled or not.

* Accessing data files through `TEST_SRCDIR` directly is error-prone
under bzlmod because the directory name within runfiles changes from the
workspace name to `_main`. To fix, use the runfiles libraries, which
know how to map apparent repo names to the actual directory name. In the
integration tests, the runfiles library isn't available, so just check
for the `_main` directory instead.

Work towards #1469
diff --git a/examples/wheel/BUILD.bazel b/examples/wheel/BUILD.bazel
index 81422d3..ab4f3a3 100644
--- a/examples/wheel/BUILD.bazel
+++ b/examples/wheel/BUILD.bazel
@@ -323,4 +323,7 @@
         ":python_requires_in_a_package",
         ":use_rule_with_dir_in_outs",
     ],
+    deps = [
+        "//python/runfiles",
+    ],
 )
diff --git a/examples/wheel/wheel_test.py b/examples/wheel/wheel_test.py
index 671bd8a..8c0f53e 100644
--- a/examples/wheel/wheel_test.py
+++ b/examples/wheel/wheel_test.py
@@ -18,18 +18,33 @@
 import unittest
 import zipfile
 
+from python.runfiles import runfiles
+
 
 class WheelTest(unittest.TestCase):
     maxDiff = None
 
+    def setUp(self):
+        super().setUp()
+        self.runfiles = runfiles.Create()
+
+    def _get_path(self, filename):
+        runfiles_path = os.path.join("rules_python/examples/wheel", filename)
+        path = self.runfiles.Rlocation(runfiles_path)
+        # The runfiles API can return None if the path doesn't exist or
+        # can't be resolved.
+        if not path:
+            raise AssertionError(f"Runfiles failed to resolve {runfiles_path}")
+        elif not os.path.exists(path):
+            # A non-None value doesn't mean the file actually exists, though
+            raise AssertionError(
+                f"Path {path} does not exist (from runfiles path {runfiles_path}"
+            )
+        else:
+            return path
+
     def test_py_library_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
-            "example_minimal_library-0.0.1-py3-none-any.whl",
-        )
+        filename = self._get_path("example_minimal_library-0.0.1-py3-none-any.whl")
         with zipfile.ZipFile(filename) as zf:
             self.assertEqual(
                 zf.namelist(),
@@ -43,11 +58,7 @@
             )
 
     def test_py_package_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "example_minimal_package-0.0.1-py3-none-any.whl",
         )
         with zipfile.ZipFile(filename) as zf:
@@ -65,11 +76,7 @@
             )
 
     def test_customized_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "example_customized-0.0.1-py3-none-any.whl",
         )
         with zipfile.ZipFile(filename) as zf:
@@ -154,31 +161,27 @@
             )
 
     def test_legacy_filename_escaping(self):
-        filename = os.path.join(
-            os.environ['TEST_SRCDIR'],
-            'rules_python',
-            'examples',
-            'wheel',
-            'file_name_escaping-0.0.1_r7-py3-none-any.whl',
+        filename = self._get_path(
+            "file_name_escaping-0.0.1_r7-py3-none-any.whl",
         )
         with zipfile.ZipFile(filename) as zf:
             self.assertEquals(
                 zf.namelist(),
                 [
-                    'examples/wheel/lib/data.txt',
-                    'examples/wheel/lib/module_with_data.py',
-                    'examples/wheel/lib/simple_module.py',
-                    'examples/wheel/main.py',
+                    "examples/wheel/lib/data.txt",
+                    "examples/wheel/lib/module_with_data.py",
+                    "examples/wheel/lib/simple_module.py",
+                    "examples/wheel/main.py",
                     # PEP calls for replacing only in the archive filename.
                     # Alas setuptools also escapes in the dist-info directory
                     # name, so let's be compatible.
-                    'file_name_escaping-0.0.1_r7.dist-info/WHEEL',
-                    'file_name_escaping-0.0.1_r7.dist-info/METADATA',
-                    'file_name_escaping-0.0.1_r7.dist-info/RECORD',
+                    "file_name_escaping-0.0.1_r7.dist-info/WHEEL",
+                    "file_name_escaping-0.0.1_r7.dist-info/METADATA",
+                    "file_name_escaping-0.0.1_r7.dist-info/RECORD",
                 ],
             )
             metadata_contents = zf.read(
-                'file_name_escaping-0.0.1_r7.dist-info/METADATA'
+                "file_name_escaping-0.0.1_r7.dist-info/METADATA"
             )
             self.assertEquals(
                 metadata_contents,
@@ -192,11 +195,7 @@
             )
 
     def test_filename_escaping(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "file_name_escaping-0.0.1rc1+ubuntu.r7-py3-none-any.whl",
         )
         with zipfile.ZipFile(filename) as zf:
@@ -230,11 +229,7 @@
             )
 
     def test_custom_package_root_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "examples_custom_package_root-0.0.1-py3-none-any.whl",
         )
 
@@ -262,11 +257,7 @@
                 self.assertFalse(line.startswith("/"))
 
     def test_custom_package_root_multi_prefix_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "example_custom_package_root_multi_prefix-0.0.1-py3-none-any.whl",
         )
 
@@ -293,11 +284,7 @@
                 self.assertFalse(line.startswith("/"))
 
     def test_custom_package_root_multi_prefix_reverse_order_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "example_custom_package_root_multi_prefix_reverse_order-0.0.1-py3-none-any.whl",
         )
 
@@ -324,11 +311,7 @@
                 self.assertFalse(line.startswith("/"))
 
     def test_python_requires_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "example_python_requires_in_a_package-0.0.1-py3-none-any.whl",
         )
         with zipfile.ZipFile(filename) as zf:
@@ -359,11 +342,7 @@
             "Windows": "win",
         }
         os_string = os_strings[platform.system()]
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             f"example_python_abi3_binary_wheel-0.0.1-cp38-abi3-{os_string}_{arch}.whl",
         )
         with zipfile.ZipFile(filename) as zf:
@@ -396,11 +375,7 @@
             )
 
     def test_rule_creates_directory_and_is_included_in_wheel(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
+        filename = self._get_path(
             "use_rule_with_dir_in_outs-0.0.1-py3-none-any.whl",
         )
 
@@ -417,12 +392,8 @@
             )
 
     def test_rule_expands_workspace_status_keys_in_wheel_metadata(self):
-        filename = os.path.join(
-            os.environ["TEST_SRCDIR"],
-            "rules_python",
-            "examples",
-            "wheel",
-            "example_minimal_library_BUILD_USER_-0.1._BUILD_TIMESTAMP_-py3-none-any.whl",
+        filename = self._get_path(
+            "example_minimal_library_BUILD_USER_-0.1._BUILD_TIMESTAMP_-py3-none-any.whl"
         )
 
         with zipfile.ZipFile(filename) as zf:
diff --git a/python/tests/toolchains/run_acceptance_test.py.tmpl b/python/tests/toolchains/run_acceptance_test.py.tmpl
index 150e1a9..5748047 100644
--- a/python/tests/toolchains/run_acceptance_test.py.tmpl
+++ b/python/tests/toolchains/run_acceptance_test.py.tmpl
@@ -16,12 +16,17 @@
 import subprocess
 import unittest
 
-
 class TestPythonVersion(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         os.chdir("%test_location%")
-        rules_python_path = os.path.join(os.environ["TEST_SRCDIR"], "rules_python")
+        test_srcdir = os.environ["TEST_SRCDIR"]
+        # When bzlmod is enabled, the name of the directory in runfiles changes
+        # to _main instead of rules_python
+        if os.path.exists(os.path.join(test_srcdir, "_main")):
+          rules_python_path = os.path.join(test_srcdir, "_main")
+        else:
+          rules_python_path = os.path.join(test_srcdir, "rules_python")
 
         test_tmpdir = os.environ["TEST_TMPDIR"]
         if %is_windows%:
@@ -57,13 +62,13 @@
 
     def test_match_toolchain(self):
         output = subprocess.check_output(
-          f"bazel run @python//:python3 -- --version",
+          f"bazel run --announce_rc @python//:python3 -- --version",
           shell = True, # Shell needed to look up via PATH
           text=True,
         ).strip()
         self.assertEqual(output, "Python %python_version%")
 
-        subprocess.run("bazel test //...", shell=True, check=True)
+        subprocess.run("bazel test --announce_rc //...", shell=True, check=True)
 
 
 if __name__ == "__main__":
diff --git a/tests/runfiles/BUILD.bazel b/tests/runfiles/BUILD.bazel
index d62e179..6193ee9 100644
--- a/tests/runfiles/BUILD.bazel
+++ b/tests/runfiles/BUILD.bazel
@@ -1,7 +1,11 @@
-load("@rules_python//python:defs.bzl", "py_test")
+load("@rules_python//python:py_test.bzl", "py_test")
+load("@rules_python//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
 
 py_test(
     name = "runfiles_test",
     srcs = ["runfiles_test.py"],
+    env = {
+        "BZLMOD_ENABLED": "1" if BZLMOD_ENABLED else "0",
+    },
     deps = ["//python/runfiles"],
 )
diff --git a/tests/runfiles/runfiles_test.py b/tests/runfiles/runfiles_test.py
index 3a1f492..5cc9568 100644
--- a/tests/runfiles/runfiles_test.py
+++ b/tests/runfiles/runfiles_test.py
@@ -514,18 +514,14 @@
             )
 
     def testCurrentRepository(self):
-        # This test assumes that it is running without --enable_bzlmod as the
-        # correct result with Bzlmod would be the empty string - the canonical
-        # name # of the main repository. Without Bzlmod, the main repository is
-        # treated just like any other repository and has the name of its
-        # runfiles directory returned, which coincides with the name specified
-        # in the WORKSPACE file.
-        #
-        # Specify a fake runfiles directory to verify that its value isn't used
-        # by the function.
+        # Under bzlmod, the current repository name is the empty string instead
+        # of the name in the workspace file.
+        if bool(int(os.environ["BZLMOD_ENABLED"])):
+            expected = ""
+        else:
+            expected = "rules_python"
         self.assertEqual(
-            runfiles.Create({"RUNFILES_DIR": "whatever"}).CurrentRepository(),
-            "rules_python",
+            runfiles.Create({"RUNFILES_DIR": "whatever"}).CurrentRepository(), expected
         )
 
     @staticmethod