Properly handle source roots and resources with V2 Pytest runner (#8063)

### Problem
While #7696 added support for source roots, its implementation of stripping the actual prefix from the file broke use of resources, specifically when referring to their path in source code, e.g. calling `open('tests/python/pants_test/backend/jvm/tasks/ivy_utils.py')`. This is because we physically stripped the name, so the hardcoded file path would need to be updated to `'pants_test/backend/jvm/tasks/ivy_utils.py'`. Per #8060 (comment), this is not acceptable because it changes a public API.

We must have a way to relativize Python files so that their import statements work, but to also keep the full path.

### Solution
By populating the `PYTHONPATH` env var for the Pex that runs Pytest with the unique source root paths, Python absolute imports work without actually having to strip any prefix. This appears similar to what [Pycharm and IntelliJ do](https://stackoverflow.com/questions/21199382/why-does-pycharm-always-add-contents-root-to-my-pythonpath-and-what-pythonpat).

#### Alternative attempted: use `PEX`'s `--source-directory`
PEX is able to understand source roots via its `--source-directory` and `--resource-directory` args. The idea was to create two pexes: a requirements PEX and a sources PEX, then to merge the two via `PEX_PATH`. They would stay separate to increase the likelihood of cache hits.

This solution failed, however, due to Pytest failing to discover any tests. See #8063 (comment) and the comment below it.

#### Alternative attempted: only strip `PythonTarget`s
The first idea was to only strip any subclass of `PythonTarget` and keep the full path for other target types like `files` (2b68ebd). This failed quickly, though, because the `VERSION` file is not able to correctly imported in this case.

### Result
Two tests can be removed from the V2 blacklist without any changes to them required.

build-support/unit_test_v2_blacklist.txt

@@ -1,9 +1,7 @@
 tests/python/pants_test/backend/codegen/antlr/java:java
 tests/python/pants_test/backend/graph_info/tasks:list_targets
-tests/python/pants_test/backend/jvm/tasks/reports:reports
 tests/python/pants_test/backend/jvm/tasks:consolidate_classpath
 tests/python/pants_test/backend/jvm/tasks:coursier_resolve
-tests/python/pants_test/backend/jvm/tasks:ivy_utils
 tests/python/pants_test/backend/jvm/tasks:scalafmt
 tests/python/pants_test/backend/jvm/tasks:scalastyle
 tests/python/pants_test/backend/python/tasks:pytest_run

pants.remote.ini

@@ -23,10 +23,13 @@ remote_execution_extra_platform_properties: [
     "container-image=docker://gcr.io/pants-remoting-beta/rbe-remote-execution@sha256:20046a8b909cbbd94f3ff33808ceeeaa489a05eea5d00ad57c378f81f21462fc",
   ]
 
-# This should correspond to the number of workers running in Google RBE.
+# This should correspond to the number of workers running in Google RBE. See
+# https://console.cloud.google.com/apis/api/remotebuildexecution.googleapis.com/quotas?project=pants-remoting-beta&folder&organizationId&duration=PT6H.
 process_execution_remote_parallelism: 8
 
 [python-setup]
+# TODO(#7735): This config is not ideal, that we must specify the PATH for both local and remote
+# platforms. This should be replaced by a proper mechanism to differentiate between the two.
 interpreter_search_paths: [
      # We include the host PATH and PEXRC values so that speculation still works.
      '<PATH>',

src/python/pants/backend/python/rules/python_test_runner.py

@@ -1,13 +1,15 @@
 # Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from typing import List, Set
+
 from pants.backend.python.rules.create_requirements_pex import (RequirementsPex,
                                                                 RequirementsPexRequest)
 from pants.backend.python.rules.inject_init import InjectedInitDigest
 from pants.backend.python.subsystems.pytest import PyTest
 from pants.backend.python.subsystems.python_setup import PythonSetup
 from pants.backend.python.subsystems.subprocess_environment import SubprocessEncodingEnvironment
-from pants.engine.fs import Digest, DirectoriesToMerge, DirectoryWithPrefixToStrip
+from pants.engine.fs import Digest, DirectoriesToMerge
 from pants.engine.isolated_process import ExecuteProcessRequest, FallibleExecuteProcessResult
 from pants.engine.legacy.graph import BuildFileAddresses, TransitiveHydratedTargets
 from pants.engine.legacy.structs import PythonTestsAdaptor
@@ -18,7 +20,6 @@
 from pants.util.strutil import create_path_env_var
 
 
-# TODO: Support resources
 # TODO(7697): Use a dedicated rule for removing the source root prefix, so that this rule
 # does not have to depend on SourceRootConfig.
 @rule(TestResult, [PythonTestsAdaptor, PyTest, PythonSetup, SourceRootConfig, SubprocessEncodingEnvironment])
@@ -61,32 +62,22 @@ def run_python_test(test_target, pytest, python_setup, source_root_config, subpr
     )
   )
 
-  source_roots = source_root_config.get_source_roots()
-
-  # Gather sources and adjust for the source root.
+  # Gather sources and source roots.
   # TODO: make TargetAdaptor return a 'sources' field with an empty snapshot instead of raising to
   # simplify the hasattr() checks here!
-  # TODO(7714): restore the full source name for the stdout of the Pytest run.
-  sources_snapshots_and_source_roots = []
+  source_roots = source_root_config.get_source_roots()
+  sources_digest_list: List[Digest] = []
+  source_root_paths: Set[str] = set()
   for maybe_source_target in all_targets:
-    if hasattr(maybe_source_target, 'sources'):
-      tgt_snapshot = maybe_source_target.sources.snapshot
-      tgt_source_root = source_roots.find_by_path(maybe_source_target.address.spec_path)
-      sources_snapshots_and_source_roots.append((tgt_snapshot, tgt_source_root))
-  all_sources_digests = yield [
-    Get(
-      Digest,
-      DirectoryWithPrefixToStrip(
-        directory_digest=snapshot.directory_digest,
-        prefix=source_root.path
-      )
-    )
-    for snapshot, source_root
-    in sources_snapshots_and_source_roots
-  ]
+    if not hasattr(maybe_source_target, 'sources'):
+      continue
+    sources_digest_list.append(maybe_source_target.sources.snapshot.directory_digest)
+    source_root = source_roots.find_by_path(maybe_source_target.address.spec_path)
+    if source_root is not None:
+      source_root_paths.add(source_root.path)
 
   sources_digest = yield Get(
-    Digest, DirectoriesToMerge(directories=tuple(all_sources_digests)),
+    Digest, DirectoriesToMerge(directories=tuple(sources_digest_list)),
   )
 
   inits_digest = yield Get(InjectedInitDigest, Digest, sources_digest)
@@ -105,6 +96,9 @@ def run_python_test(test_target, pytest, python_setup, source_root_config, subpr
   interpreter_search_paths = create_path_env_var(python_setup.interpreter_search_paths)
   pex_exe_env = {
     'PATH': interpreter_search_paths,
+    # This adds support for absolute imports of source roots, e.g.
+    # `src/python/pants` -> `import pants`.
+    'PYTHONPATH': create_path_env_var(sorted(source_root_paths)),
     **subprocess_encoding_environment.invocation_environment_dict
   }
 

src/python/pants/engine/rules.py

@@ -30,7 +30,7 @@ class _RuleVisitor(ast.NodeVisitor):
   """Pull `Get` calls out of an @rule body and validate `yield` statements."""
 
   def __init__(self, func, func_node, func_source, orig_indent, parents_table):
-    super(_RuleVisitor, self).__init__()
+    super().__init__()
     self._gets = []
     self._func = func
     self._func_node = func_node
@@ -218,11 +218,13 @@ def wrapper(func):
     def resolve_type(name):
       resolved = getattr(owning_module, name, None) or owning_module.__builtins__.get(name, None)
       if resolved is None:
-        raise ValueError('Could not resolve type `{}` in top level of module {}'
-                         .format(name, owning_module.__name__))
+        raise ValueError(
+          f'Could not resolve type `{name}` in top level of module {owning_module.__name__}'
+        )
       elif not isinstance(resolved, type):
-        raise ValueError('Expected a `type` constructor for `{}`, but got: {} (type `{}`)'
-                         .format(name, resolved, type(resolved).__name__))
+        raise ValueError(
+          f'Expected a `type` constructor for `{name}`, but got: {resolved} (type `{type(resolved).__name__}`)'
+        )
       return resolved
 
     gets = OrderedSet()

tests/python/pants_test/backend/python/rules/BUILD

@@ -36,5 +36,6 @@ python_tests(
     'tests/python/pants_test:int-test',
   ],
   tags={'integration'},
+  timeout = 90,
 )
 

tests/python/pants_test/backend/python/rules/test_python_test_runner_integration.py

@@ -65,7 +65,7 @@ def test_passing_test(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_pass.py .                                             [100%]
+        testprojects/tests/python/pants/dummies/test_pass.py .                   [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -88,7 +88,7 @@ def test_failing_test(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_fail.py F                                             [100%]
+        testprojects/tests/python/pants/dummies/test_fail.py F                   [100%]
 
         =================================== FAILURES ===================================
         __________________________________ test_fail ___________________________________
@@ -97,7 +97,7 @@ def test_fail():
         >     assert False
         E     assert False
 
-        pants/dummies/test_fail.py:2: AssertionError
+        testprojects/tests/python/pants/dummies/test_fail.py:2: AssertionError
         =========================== 1 failed in SOME_TEXT ===========================
 
 
@@ -120,7 +120,7 @@ def test_mixed_tests(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_fail.py F                                             [100%]
+        testprojects/tests/python/pants/dummies/test_fail.py F                   [100%]
 
         =================================== FAILURES ===================================
         __________________________________ test_fail ___________________________________
@@ -129,7 +129,7 @@ def test_fail():
         >     assert False
         E     assert False
 
-        pants/dummies/test_fail.py:2: AssertionError
+        testprojects/tests/python/pants/dummies/test_fail.py:2: AssertionError
         =========================== 1 failed in SOME_TEXT ===========================
 
         testprojects/tests/python/pants/dummies:passing_target stdout:
@@ -139,7 +139,7 @@ def test_fail():
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_pass.py .                                             [100%]
+        testprojects/tests/python/pants/dummies/test_pass.py .                   [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -163,7 +163,7 @@ def test_source_dep_absolute_import(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_with_source_dep_absolute_import.py .                  [100%]
+        testprojects/tests/python/pants/dummies/test_with_source_dep_absolute_import.py . [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -186,7 +186,7 @@ def test_source_dep_relative_import(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_with_source_dep_relative_import.py .                  [100%]
+        testprojects/tests/python/pants/dummies/test_with_source_dep_relative_import.py . [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -209,7 +209,7 @@ def test_thirdparty_dep(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_with_thirdparty_dep.py .                              [100%]
+        testprojects/tests/python/pants/dummies/test_with_thirdparty_dep.py .    [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -232,7 +232,7 @@ def test_transitive_dep(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        pants/dummies/test_with_transitive_dep.py .                              [100%]
+        testprojects/tests/python/pants/dummies/test_with_transitive_dep.py .    [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 

tests/python/pants_test/engine/test_engine.py

@@ -29,7 +29,7 @@ class D:
 
 
 def fn_raises(x):
-  raise Exception('An exception for {}'.format(type(x).__name__))
+  raise Exception(f'An exception for {type(x).__name__}')
 
 
 @rule(A, [B])
@@ -119,16 +119,16 @@ def test_include_trace_error_raises_error_with_trace(self):
 
     self.assert_equal_with_printing(dedent('''
       1 Exception encountered:
-      Computing Select(<pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A)
-        Computing Task(nested_raise(), <pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A, true)
+      Computing Select(<tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A)
+        Computing Task(nested_raise(), <tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A, true)
           Throw(An exception for B)
             Traceback (most recent call last):
               File LOCATION-INFO, in call
                 val = func(*args)
               File LOCATION-INFO, in nested_raise
                 fn_raises(x)
               File LOCATION-INFO, in fn_raises
-                raise Exception('An exception for {}'.format(type(x).__name__))
+                raise Exception(f'An exception for {type(x).__name__}')
             Exception: An exception for B
       ''').lstrip()+'\n',
       remove_locations_from_traceback(str(cm.exception)))
@@ -171,9 +171,9 @@ def a_from_c_and_d(c, d):
 
     self.assert_equal_with_printing(dedent('''
       1 Exception encountered:
-      Computing Select(<pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A)
-        Computing Task(a_from_c_and_d(), <pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A, true)
-          Computing Task(d_from_b_nested_raise(), <pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, =D, true)
+      Computing Select(<tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A)
+        Computing Task(a_from_c_and_d(), <tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A, true)
+          Computing Task(d_from_b_nested_raise(), <tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, =D, true)
             Throw(An exception for B)
               Traceback (most recent call last):
                 File LOCATION-INFO, in call
@@ -185,9 +185,9 @@ def a_from_c_and_d(c, d):
               Exception: An exception for B
 
 
-      Computing Select(<pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A)
-        Computing Task(a_from_c_and_d(), <pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A, true)
-          Computing Task(c_from_b_nested_raise(), <pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, =C, true)
+      Computing Select(<tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A)
+        Computing Task(a_from_c_and_d(), <tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, A, true)
+          Computing Task(c_from_b_nested_raise(), <tests.python.pants_test.engine.test_engine.B object at 0xEEEEEEEEE>, =C, true)
             Throw(An exception for B)
               Traceback (most recent call last):
                 File LOCATION-INFO, in call