Add support for source roots in V2 ./pants test

src/python/pants/backend/python/rules/BUILD

@@ -7,5 +7,6 @@ python_library(
     'src/python/pants/engine:rules',
     'src/python/pants/engine:selectors',
     'src/python/pants/rules/core',
+    'src/python/pants/source:source',
   ],
 )

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

@@ -10,15 +10,18 @@
 
 from future.utils import text_type
 
+from pants.backend.python.subsystems.pex_build_util import identify_missing_init_files
 from pants.backend.python.subsystems.pytest import PyTest
 from pants.backend.python.subsystems.python_setup import PythonSetup
-from pants.engine.fs import Digest, MergedDirectories, Snapshot, UrlToFetch
+from pants.engine.fs import (Digest, DirectoryWithPrefixToStrip, FilesContent, MergedDirectories,
+                             Snapshot, UrlToFetch)
 from pants.engine.isolated_process import (ExecuteProcessRequest, ExecuteProcessResult,
                                            FallibleExecuteProcessResult)
 from pants.engine.legacy.graph import TransitiveHydratedTarget
 from pants.engine.rules import optionable_rule, rule
 from pants.engine.selectors import Get
 from pants.rules.core.core_test_model import Status, TestResult
+from pants.source.source_root import SourceRootConfig
 
 
 # This class currently exists so that other rules could be added which turned a HydratedTarget into
@@ -44,8 +47,12 @@ def parse_interpreter_constraints(python_setup, python_target_adaptors):
 
 # TODO: Support deps
 # TODO: Support resources
-@rule(PyTestResult, [TransitiveHydratedTarget, PyTest, PythonSetup])
-def run_python_test(transitive_hydrated_target, pytest, python_setup):
+# TODO: this rule should not directly require SourceRootConfig. Arguably, we should
+# have a dedicated rule that will allow us to do something like this:
+# `yield Get(SourceRootDigest, Digest, digest)`. To see an example of what this might
+# look like, see https://github.com/pantsbuild/pants/blob/f25c5b9ec34cea8bae4e6ea47cbec5caefd81576/src/python/pants/rules/core/source_roots.py.
+@rule(PyTestResult, [TransitiveHydratedTarget, PyTest, PythonSetup, SourceRootConfig])
+def run_python_test(transitive_hydrated_target, pytest, python_setup, source_root_config):
   target_root = transitive_hydrated_target.root
 
   # TODO: Inject versions and digests here through some option, rather than hard-coding it.
@@ -98,17 +105,53 @@ def run_python_test(transitive_hydrated_target, pytest, python_setup):
   requirements_pex_response = yield Get(
     ExecuteProcessResult, ExecuteProcessRequest, requirements_pex_request)
 
-  # Gather sources.
+  source_roots = source_root_config.get_source_roots()
+
+  # Gather sources and adjust for the source root.
   # TODO: make TargetAdaptor return a 'sources' field with an empty snapshot instead of raising to
   # simplify the hasattr() checks here!
+  # TODO: find some way to restore the full source name for the stdout of the Pytest run, so that
+  # we can output the original path, rather than the path without the source root. It appears we
+  # do this in V1 via parsing the Junit XML file and adding the value back? See pytest_run.py's
+  # _get_target_from_test(). This is a low priority.
   all_sources_digests = []
   for maybe_source_target in all_targets:
     if hasattr(maybe_source_target.adaptor, 'sources'):
       sources_snapshot = maybe_source_target.adaptor.sources.snapshot
-      all_sources_digests.append(sources_snapshot.directory_digest)
+      digest_adjusted_for_source_root = yield Get(
+        Digest,
+        DirectoryWithPrefixToStrip(
+          sources_snapshot.directory_digest,
+          source_roots.find_by_path(maybe_source_target.adaptor.address.spec_path).path
+        )
+      )
+      all_sources_digests.append(digest_adjusted_for_source_root)
+
+  sources_digest = yield Get(
+    Digest, MergedDirectories(directories=tuple(all_sources_digests)),
+  )
+
+  # TODO: add intrinsic to go from Digest->Snapshot, so that we can avoid having to use
+  # `FileContent`, which unnecessarily gets the content.
+  # TODO: somehow generalize this injection of __init__.py files. Presumably it will be useful
+  # for other Python commands.
+  file_contents = yield Get(FilesContent, Digest, sources_digest)
+  file_paths = [fc.path for fc in file_contents]
+  injected_inits = tuple(sorted(identify_missing_init_files(file_paths)))
+  touch_init_request = ExecuteProcessRequest(
+    argv=("touch",) + injected_inits,
+    output_files=injected_inits,
+    description="Inject empty __init__.py into all packages without one already.",
+    input_files=sources_digest,
+    env={"PATH": os.environ["PATH"]}
+  )
+
+  touch_init_result = yield Get(ExecuteProcessResult, ExecuteProcessRequest, touch_init_request)
 
-  all_input_digests = all_sources_digests + [
+  all_input_digests = [
+    sources_digest,
     requirements_pex_response.output_directory_digest,
+    touch_init_result.output_directory_digest
   ]
   merged_input_files = yield Get(
     Digest,
@@ -138,4 +181,5 @@ def rules():
       run_python_test,
       optionable_rule(PyTest),
       optionable_rule(PythonSetup),
+      optionable_rule(SourceRootConfig),
     ]

src/python/pants/backend/python/subsystems/pex_build_util.py

@@ -78,6 +78,27 @@ def targets_by_platform(targets, python_setup):
   return d
 
 
+def identify_missing_init_files(sources):
+  """Return the list of paths that would need to be added to ensure that every package has
+  an __init__.py. """
+  packages = set()
+  for source in sources:
+    if source.endswith('.py'):
+      pkg_dir = os.path.dirname(source)
+      if pkg_dir and pkg_dir not in packages:
+        package = ''
+        for component in pkg_dir.split(os.sep):
+          package = os.path.join(package, component)
+          packages.add(package)
+
+  missing_init_files = set()
+  for package in packages:
+    init_file = os.path.join(package, '__init__.py')
+    if init_file not in sources:
+      missing_init_files.add(init_file)
+  return missing_init_files
+
+
 def _create_source_dumper(builder, tgt):
   if type(tgt) == Files:
     # Loose `Files` as opposed to `Resources` or `PythonTarget`s have no (implied) package structure
@@ -305,31 +326,14 @@ def add_sources_from(self, tgt):
   def _prepare_inits(self):
     chroot = self._builder.chroot()
     sources = chroot.get('source') | chroot.get('resource')
-
-    packages = set()
-    for source in sources:
-      if source.endswith('.py'):
-        pkg_dir = os.path.dirname(source)
-        if pkg_dir and pkg_dir not in packages:
-          package = ''
-          for component in pkg_dir.split(os.sep):
-            package = os.path.join(package, component)
-            packages.add(package)
-
-    missing_pkg_files = []
-    for package in packages:
-      pkg_file = os.path.join(package, '__init__.py')
-      if pkg_file not in sources:
-        missing_pkg_files.append(pkg_file)
-
-    if missing_pkg_files:
+    missing_init_files = identify_missing_init_files(sources)
+    if missing_init_files:
       with temporary_file() as ns_package:
         ns_package.write(b'__import__("pkg_resources").declare_namespace(__name__)')
         ns_package.flush()
-        for missing_pkg_file in missing_pkg_files:
-          self._builder.add_source(ns_package.name, missing_pkg_file)
-
-    return missing_pkg_files
+        for missing_init_file in missing_init_files:
+          self._builder.add_source(ns_package.name, missing_init_file)
+    return missing_init_files
 
   def set_emit_warnings(self, emit_warnings):
     self._builder.info.emit_warnings = emit_warnings

testprojects/tests/python/pants/dummies/BUILD

@@ -10,15 +10,20 @@ python_tests(
 
 python_library(
   name = 'example_lib',
-  sources = [
-    'example_source.py',
-    '__init__.py',
+  sources = ['example_source.py'],
+)
+
+python_tests(
+  name = 'target_with_source_dep_absolute_import',
+  sources = ['test_with_source_dep_absolute_import.py'],
+  dependencies = [
+    ':example_lib',
   ],
 )
 
 python_tests(
-  name = 'target_with_source_dep',
-  sources = ['test_with_source_dep.py'],
+  name = 'target_with_source_dep_relative_import',
+  sources = ['test_with_source_dep_relative_import.py'],
   dependencies = [
     ':example_lib',
   ],

testprojects/tests/python/pants/dummies/test_with_source_dep_absolute_import.py

@@ -0,0 +1,7 @@
+from __future__ import absolute_import
+
+from pants.dummies.example_source import add_two
+
+
+def test_external_method():
+  assert add_two(2) == 4

tests/python/pants_test/rules/test_test_integration.py

@@ -72,7 +72,7 @@ def test_passing_python_test(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        testprojects/tests/python/pants/dummies/test_pass.py .                   [100%]
+        pants/dummies/test_pass.py .                                             [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -95,7 +95,7 @@ def test_failing_python_test(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        testprojects/tests/python/pants/dummies/test_fail.py F                   [100%]
+        pants/dummies/test_fail.py F                                             [100%]
 
         =================================== FAILURES ===================================
         __________________________________ test_fail ___________________________________
@@ -104,36 +104,59 @@ def test_fail():
         >     assert False
         E     assert False
 
-        testprojects/tests/python/pants/dummies/test_fail.py:2: AssertionError
+        pants/dummies/test_fail.py:2: AssertionError
         =========================== 1 failed in SOME_TEXT ===========================
 
 
         testprojects/tests/python/pants/dummies:failing_target                          .....   FAILURE
         """),
     )
 
-  def test_source_dep(self):
+  def test_source_dep_absolute_import(self):
     pants_run = self.run_passing_pants_test([
-      'testprojects/tests/python/pants/dummies:target_with_source_dep',
+      'testprojects/tests/python/pants/dummies:target_with_source_dep_absolute_import',
     ])
     self.assert_fuzzy_string_match(
       pants_run.stdout_data,
       dedent("""\
-        testprojects/tests/python/pants/dummies:target_with_source_dep stdout:
+        testprojects/tests/python/pants/dummies:target_with_source_dep_absolute_import stdout:
         ============================= test session starts ==============================
         platform SOME_TEXT
         rootdir: SOME_TEXT
         plugins: SOME_TEXT
         collected 1 item
 
-        testprojects/tests/python/pants/dummies/test_with_source_dep.py .        [100%]
+        pants/dummies/test_with_source_dep_absolute_import.py .                  [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
 
-        testprojects/tests/python/pants/dummies:target_with_source_dep                  .....   SUCCESS
+        testprojects/tests/python/pants/dummies:target_with_source_dep_absolute_import  .....   SUCCESS
         """)
-    )
+      )
+
+  def test_source_dep_relative_import(self):
+    pants_run = self.run_passing_pants_test([
+      'testprojects/tests/python/pants/dummies:target_with_source_dep_relative_import',
+    ])
+    self.assert_fuzzy_string_match(
+      pants_run.stdout_data,
+      dedent("""\
+        testprojects/tests/python/pants/dummies:target_with_source_dep_relative_import stdout:
+        ============================= test session starts ==============================
+        platform SOME_TEXT
+        rootdir: SOME_TEXT
+        plugins: SOME_TEXT
+        collected 1 item
+
+        pants/dummies/test_with_source_dep_relative_import.py .                  [100%]
+
+        =========================== 1 passed in SOME_TEXT ===========================
+
+
+        testprojects/tests/python/pants/dummies:target_with_source_dep_relative_import  .....   SUCCESS
+        """)
+      )
 
   def test_thirdparty_dep(self):
     pants_run = self.run_passing_pants_test([
@@ -149,7 +172,7 @@ def test_thirdparty_dep(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        testprojects/tests/python/pants/dummies/test_with_thirdparty_dep.py .    [100%]
+        pants/dummies/test_with_thirdparty_dep.py .                              [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================
 
@@ -173,7 +196,7 @@ def test_mixed_python_tests(self):
         plugins: SOME_TEXT
         collected 1 item
 
-        testprojects/tests/python/pants/dummies/test_fail.py F                   [100%]
+        pants/dummies/test_fail.py F                                             [100%]
 
         =================================== FAILURES ===================================
         __________________________________ test_fail ___________________________________
@@ -182,7 +205,7 @@ def test_fail():
         >     assert False
         E     assert False
 
-        testprojects/tests/python/pants/dummies/test_fail.py:2: AssertionError
+        pants/dummies/test_fail.py:2: AssertionError
         =========================== 1 failed in SOME_TEXT ===========================
 
         testprojects/tests/python/pants/dummies:passing_target stdout:
@@ -192,7 +215,7 @@ def test_fail():
         plugins: SOME_TEXT
         collected 1 item
 
-        testprojects/tests/python/pants/dummies/test_pass.py .                   [100%]
+        pants/dummies/test_pass.py .                                             [100%]
 
         =========================== 1 passed in SOME_TEXT ===========================