Repository: pantsbuild/pants

Issue:
Consider factoring up shared partitioning & caching/invalidation code from JUnitRun and PytestRun
Right now both of these tasks partition, invalidate & cache and report in very similar manners. Much code could be shared.


Discussion:

Resolving patch:
diff --git a/src/python/pants/backend/jvm/tasks/junit_run.py b/src/python/pants/backend/jvm/tasks/junit_run.py
--- a/src/python/pants/backend/jvm/tasks/junit_run.py
+++ b/src/python/pants/backend/jvm/tasks/junit_run.py
@@ -29,17 +29,16 @@
 from pants.backend.jvm.tasks.reports.junit_html_report import JUnitHtmlReport, NoJunitHtmlReport
 from pants.base.build_environment import get_buildroot
 from pants.base.deprecated import deprecated_conditional
-from pants.base.exceptions import ErrorWhileTesting, TargetDefinitionException, TaskError
+from pants.base.exceptions import TargetDefinitionException, TaskError
 from pants.base.workunit import WorkUnitLabel
 from pants.build_graph.files import Files
 from pants.build_graph.target import Target
 from pants.build_graph.target_scopes import Scopes
-from pants.invalidation.cache_manager import VersionedTargetSet
 from pants.java.distribution.distribution import DistributionLocator
 from pants.java.executor import SubprocessExecutor
 from pants.java.junit.junit_xml_parser import RegistryOfTests, Test, parse_failed_targets
 from pants.process.lock import OwnerPrintingInterProcessFileLock
-from pants.task.testrunner_task_mixin import TestResult, TestRunnerTaskMixin
+from pants.task.testrunner_task_mixin import PartitionedTestRunnerTaskMixin, TestResult
 from pants.util import desktop
 from pants.util.argutil import ensure_arg, remove_arg
 from pants.util.contextutil import environment_as, temporary_dir
@@ -131,7 +130,7 @@ def iter_possible_tests(self, context):
     yield Test(classname=self._classname, methodname=self._methodname)
 
 
-class JUnitRun(TestRunnerTaskMixin, JvmToolTaskMixin, JvmTask):
+class JUnitRun(PartitionedTestRunnerTaskMixin, JvmToolTaskMixin, JvmTask):
   """
   :API: public
   """
@@ -146,10 +145,6 @@ def implementation_version(cls):
   def register_options(cls, register):
     super(JUnitRun, cls).register_options(register)
 
-    register('--fast', type=bool, default=True, fingerprint=True,
-             help='Run all tests in a single junit invocation. If turned off, each test target '
-                  'will run in its own junit invocation, which will be slower, but isolates '
-                  'tests from process-wide state created by tests in other targets.')
     register('--batch-size', advanced=True, type=int, default=cls._BATCH_ALL, fingerprint=True,
              help='Run at most this many tests in a single test process.')
     register('--test', type=list, fingerprint=True,
@@ -174,11 +169,6 @@ def register_options(cls, register):
              help='Set the working directory. If no argument is passed, use the build root. '
                   'If cwd is set on a target, it will supersede this option. It is an error to '
                   'use this option in combination with `--chroot`')
-    register('--chroot', advanced=True, fingerprint=True, type=bool, default=False,
-             help='Run tests in a chroot. Any loose files tests depend on via `{}` dependencies '
-                  'will be copied to the chroot. If cwd is set on a target, it will supersede this'
-                  'option. It is an error to use this option in combination with `--cwd`'
-                  .format(Files.alias()))
     register('--strict-jvm-version', type=bool, advanced=True, fingerprint=True,
              help='If true, will strictly require running junits with the same version of java as '
                   'the platform -target level. Otherwise, the platform -target level will be '
@@ -234,13 +224,12 @@ def __init__(self, *args, **kwargs):
     options = self.get_options()
     self._tests_to_run = options.test
     self._batch_size = options.batch_size
-    self._fail_fast = options.fail_fast
 
-    if options.cwd and options.chroot:
+    if options.cwd and self.run_tests_in_chroot:
       raise self.OptionError('Cannot set both `cwd` ({}) and ask for a `chroot` at the same time.'
                              .format(options.cwd))
 
-    if options.chroot:
+    if self.run_tests_in_chroot:
       self._working_dir = None
     else:
       self._working_dir = options.cwd or get_buildroot()
@@ -252,7 +241,7 @@ def __init__(self, *args, **kwargs):
     self._legacy_report_layout = options.legacy_report_layout
 
   @memoized_method
-  def _args(self, output_dir):
+  def _args(self, fail_fast, output_dir):
     args = self.args[:]
 
     options = self.get_options()
@@ -263,7 +252,7 @@ def _args(self, output_dir):
     else:
       args.append('-output-mode=NONE')
 
-    if self._fail_fast:
+    if fail_fast:
       args.append('-fail-fast')
     args.append('-outdir')
     args.append(output_dir)
@@ -398,15 +387,11 @@ def _chroot(self, targets, workdir):
         )
         yield chroot
 
-  @property
-  def _per_target(self):
-    return not self.get_options().fast
-
   @property
   def _batched(self):
     return self._batch_size != self._BATCH_ALL
 
-  def _run_junit(self, test_targets, output_dir, coverage):
+  def run_tests(self, fail_fast, test_targets, output_dir, coverage):
     test_registry = self._collect_test_targets(test_targets)
     if test_registry.empty:
       return TestResult.rc(0)
@@ -447,7 +432,7 @@ def parse_error_handler(parse_error):
       distribution = JvmPlatform.preferred_jvm_distribution([platform], self._strict_jvm_version)
 
       # Override cmdline args with values from junit_test() target that specify concurrency:
-      args = self._args(batch_output_dir) + [u'-xmlreport']
+      args = self._args(fail_fast, batch_output_dir) + [u'-xmlreport']
 
       if concurrency is not None:
         args = remove_arg(args, '-default-parallel')
@@ -495,7 +480,7 @@ def parse_error_handler(parse_error):
           self.report_all_info_for_single_test(self.options_scope, test_target,
                                                test_name, test_info)
 
-        if result != 0 and self._fail_fast:
+        if result != 0 and fail_fast:
           break
 
     if result == 0:
@@ -573,119 +558,34 @@ def _validate_target(self, target):
       msg = 'JUnitTests target must include a non-empty set of sources.'
       raise TargetDefinitionException(target, msg)
 
-  @staticmethod
-  def _vts_for_partition(invalidation_check):
-    return VersionedTargetSet.from_versioned_targets(invalidation_check.all_vts)
-
-  def check_artifact_cache_for(self, invalidation_check):
-    # We generate artifacts, namely coverage reports, that cover the full target set.
-    return [self._vts_for_partition(invalidation_check)]
-
-  @staticmethod
-  def _collect_files(directory):
+  def collect_files(self, output_dir, coverage):
     def files_iter():
-      for dir_path, _, file_names in os.walk(directory):
+      for dir_path, _, file_names in os.walk(output_dir):
         for filename in file_names:
           yield os.path.join(dir_path, filename)
     return list(files_iter())
 
-  def _iter_partitions(self, targets, output_dir):
-    if self._per_target:
-      for target in targets:
-        yield (target,), os.path.join(output_dir, target.id)
-    else:
-      if targets:
-        yield tuple(targets), output_dir
-
-  def _execute(self, all_targets):
-    with self._isolation(all_targets) as (output_dir, reports, coverage):
-      results = {}
-      failure = False
-      for (partition, partition_output_dir) in self._iter_partitions(self._get_test_targets(),
-                                                                     output_dir):
-        try:
-          rv = self._run_partition(test_targets=partition,
-                                   output_dir=partition_output_dir,
-                                   coverage=coverage)
-        except ErrorWhileTesting as e:
-          rv = TestResult.from_error(e)
-
-        results[partition] = rv
-        if not rv.success:
-          failure = True
-          if self._fail_fast:
-            break
-
-      for partition in sorted(results):
-        rv = results[partition]
-        if len(partition) == 1 or rv.success:
-          log = self.context.log.info if rv.success else self.context.log.error
-          for target in partition:
-            log('{0:80}.....{1:>10}'.format(target.address.reference(), rv))
-        else:
-          # There is not much useful we can display in summary for a multi-target partition with
-          # failures without parsing those failures to link them to individual targets; ie: targets
-          # 2 and 8 failed in this partition of 10 targets.
-          # TODO(John Sirois): Punting here works since we have in practice just 2 partitionings:
-          # 1. All targets in singleton partitions
-          # 2. All targets in 1 partition
-          # If we get to the point where we have multiple partitions with multiple targets, some
-          # sort of summary for the multi-target partitions will probably be needed.
-          pass
-
-      msgs = [str(_rv) for _rv in results.values() if not _rv.success]
-      failed_targets = [target
-                        for _rv in results.values() if not _rv.success
-                        for target in _rv.failed_targets]
-      if len(failed_targets) > 0:
-        error = ErrorWhileTesting('\n'.join(msgs), failed_targets=failed_targets)
-      elif failure:
-        # A low-level test execution failure occurred before tests were run.
-        error = TaskError()
-      else:
-        error = None
-
-      reports.generate(output_dir, exc=error)
-      if error:
-        raise error
-
-  def _run_partition(self, test_targets, output_dir, coverage):
-    with self.invalidated(targets=test_targets,
-                          # Re-run tests when the code they test (and depend on) changes.
-                          invalidate_dependents=True) as invalidation_check:
-
-      invalid_test_tgts = [invalid_test_tgt
-                           for vts in invalidation_check.invalid_vts
-                           for invalid_test_tgt in vts.targets]
-
-      # Processing proceeds through:
-      # 1.) output -> output_dir
-      # 2.) [iff all == invalid] output_dir -> cache: We do this manually for now.
-      # 3.) [iff invalid == 0 and all > 0] cache -> workdir: Done transparently by `invalidated`.
-
-      # 1.) Write all results that will be potentially cached to output_dir.
-      result = self._run_junit(invalid_test_tgts, output_dir, coverage).checked()
-
-      cache_vts = self._vts_for_partition(invalidation_check)
-      if invalidation_check.all_vts == invalidation_check.invalid_vts:
-        # 2.) All tests in the partition were invalid, cache successful test results.
-        if result.success and self.artifact_cache_writes_enabled():
-          self.update_artifact_cache([(cache_vts, self._collect_files(output_dir))])
-      elif not invalidation_check.invalid_vts:
-        # 3.) The full partition was valid, our results will have been staged for/by caching
-        # if not already local.
-        pass
+  @contextmanager
+  def partitions(self, per_target, all_targets, test_targets):
+    with self._isolation(per_target, all_targets) as (output_dir, reports, coverage):
+      if per_target:
+        def iter_partitions():
+          for test_target in test_targets:
+            partition = (test_target,)
+            args = (os.path.join(output_dir, test_target.id), coverage)
+            yield partition, args
       else:
-        # The partition was partially invalid.
-
-        # We don't cache results; so others will need to re-run this partition.
-        # NB: We will presumably commit this change now though and so others will get this
-        # partition in a state that executes successfully; so when the 1st of the others
-        # executes against this partition; they will hit `all_vts == invalid_vts` and
-        # cache the results. That 1st of others is hopefully CI!
-        cache_vts.force_invalidate()
-
-      return result
+        def iter_partitions():
+          if test_targets:
+            partition = tuple(test_targets)
+            args = (output_dir, coverage)
+            yield partition, args
+
+      try:
+        yield iter_partitions
+      finally:
+        _, error, _ = sys.exc_info()
+        reports.generate(output_dir, exc=error)
 
   class Reports(object):
     def __init__(self, junit_html_report, coverage):
@@ -707,9 +607,9 @@ def _maybe_open_report(self, report_file_path):
           raise TaskError(e)
 
   @contextmanager
-  def _isolation(self, all_targets):
+  def _isolation(self, per_target, all_targets):
     run_dir = '_runs'
-    mode_dir = 'isolated' if self._per_target else 'combined'
+    mode_dir = 'isolated' if per_target else 'combined'
     batch_dir = str(self._batch_size) if self._batched else 'all'
     output_dir = os.path.join(self.workdir,
                               run_dir,
diff --git a/src/python/pants/backend/python/tasks/pytest_run.py b/src/python/pants/backend/python/tasks/pytest_run.py
--- a/src/python/pants/backend/python/tasks/pytest_run.py
+++ b/src/python/pants/backend/python/tasks/pytest_run.py
@@ -27,11 +27,9 @@
 from pants.base.fingerprint_strategy import DefaultFingerprintStrategy
 from pants.base.hash_utils import Sharder
 from pants.base.workunit import WorkUnitLabel
-from pants.build_graph.files import Files
 from pants.build_graph.target import Target
-from pants.invalidation.cache_manager import VersionedTargetSet
 from pants.task.task import Task
-from pants.task.testrunner_task_mixin import TestResult, TestRunnerTaskMixin
+from pants.task.testrunner_task_mixin import PartitionedTestRunnerTaskMixin, TestResult
 from pants.util.contextutil import pushd, temporary_dir, temporary_file
 from pants.util.dirutil import mergetree, safe_mkdir, safe_mkdir_for
 from pants.util.memo import memoized_method, memoized_property
@@ -89,7 +87,7 @@ def _map_exit_code(cls, value):
     return 0 if value in cls._SUCCESS_EXIT_CODES else value
 
 
-class PytestRun(TestRunnerTaskMixin, Task):
+class PytestRun(PartitionedTestRunnerTaskMixin, Task):
 
   @classmethod
   def implementation_version(cls):
@@ -98,15 +96,6 @@ def implementation_version(cls):
   @classmethod
   def register_options(cls, register):
     super(PytestRun, cls).register_options(register)
-    register('--fast', type=bool, default=True, fingerprint=True,
-             help='Run all tests in a single pytest invocation. If turned off, each test target '
-                  'will run in its own pytest invocation, which will be slower, but isolates '
-                  'tests from process-wide state created by tests in other targets.')
-
-    register('--chroot', advanced=True, fingerprint=True, type=bool, default=False,
-             help='Run tests in a chroot. Any loose files tests depend on via `{}` dependencies '
-                  'will be copied to the chroot.'
-             .format(Files.alias()))
 
     # NB: We always produce junit xml privately, and if this option is specified, we then copy
     # it to the user-specified directory, post any interaction with the cache to retrieve the
@@ -157,12 +146,6 @@ def target_filter(target):
   def _validate_target(self, target):
     pass
 
-  def _execute(self, all_targets):
-    test_targets = self._get_test_targets()
-    if test_targets:
-      self.context.release_lock()
-      self._run_tests(test_targets)
-
   class InvalidShardSpecification(TaskError):
     """Indicates an invalid `--test-shard` option."""
 
@@ -504,72 +487,34 @@ def _get_target_from_test(self, test_info, targets):
     file_info = test_info['file']
     return relsrc_to_target.get(file_info)
 
-  def _iter_partitions(self, targets):
-    # TODO(John Sirois): Consume `py.test` pexes matched to the partitioning in effect after
-    # https://github.com/pantsbuild/pants/pull/4638 lands.
-    if self.get_options().fast:
+  @contextmanager
+  def partitions(self, per_target, all_targets, test_targets):
+    if per_target:
+      def iter_partitions():
+        for test_target in test_targets:
+          yield (test_target,)
+    else:
       targets_by_target_base = OrderedDict()
-      for target in targets:
-        targets_for_base = targets_by_target_base.get(target.target_base)
+      for test_target in test_targets:
+        targets_for_base = targets_by_target_base.get(test_target.target_base)
         if targets_for_base is None:
           targets_for_base = []
-          targets_by_target_base[target.target_base] = targets_for_base
-        targets_for_base.append(target)
-      for targets in targets_by_target_base.values():
-        yield tuple(targets)
-    else:
-      for target in targets:
-        yield (target,)
+          targets_by_target_base[test_target.target_base] = targets_for_base
+        targets_for_base.append(test_target)
 
-  def _run_tests(self, targets):
-    results = {}
-    failure = False
-    for partition in self._iter_partitions(targets):
-      try:
-        rv = self._do_run_tests(partition)
-      except ErrorWhileTesting as e:
-        rv = PytestResult.from_error(e)
-      results[partition] = rv
-      if not rv.success:
-        failure = True
-        if self.get_options().fail_fast:
-          break
-
-    for partition in sorted(results):
-      rv = results[partition]
-      if len(partition) == 1 or rv.success:
-        log = self.context.log.info if rv.success else self.context.log.error
-        for target in partition:
-          log('{0:80}.....{1:>10}'.format(target.address.reference(), rv))
-      else:
-        # There is not much useful we can display in summary for a multi-target partition with
-        # failures without parsing those failures to link them to individual targets; ie: targets
-        # 2 and 8 failed in this partition of 10 targets.
-        # TODO(John Sirois): Punting here works since we have in practice just 2 partitionings:
-        # 1. All targets in singleton partitions
-        # 2. All targets in 1 partition
-        # If we get to the point where we have multiple partitions with multiple targets, some sort
-        # of summary for the multi-target partitions will probably be needed.
-        pass
-
-    failed_targets = [target
-                      for _rv in results.values() if not _rv.success
-                      for target in _rv.failed_targets]
-    if failed_targets:
-      raise ErrorWhileTesting(failed_targets=failed_targets)
-    elif failure:
-      # A low-level test execution failure occurred before tests were run.
-      raise TaskError()
+      def iter_partitions():
+        for test_targets in targets_by_target_base.values():
+          yield tuple(test_targets)
 
-  @staticmethod
-  def _vts_for_partition(invalidation_check):
-    return VersionedTargetSet.from_versioned_targets(invalidation_check.all_vts)
+    workdir = self.workdir
 
-  def check_artifact_cache_for(self, invalidation_check):
-    # We generate artifacts, namely junit.xml and coverage reports, that cover the full target set
-    # whether that is all targets in the context (`--fast`) or each target
-    # individually (`--no-fast`).
-    return [self._vts_for_partition(invalidation_check)]
+    def iter_partitions_with_args():
+      for partition in iter_partitions():
+        workdirs = _Workdirs.for_partition(workdir, partition)
+        args = (workdirs,)
+        yield partition, args
+
+    yield iter_partitions_with_args
 
   # TODO(John Sirois): Its probably worth generalizing a means to mark certain options or target
   # attributes as making results un-cacheable. See: https://github.com/pantsbuild/pants/issues/4748
@@ -577,7 +522,7 @@ class NeverCacheFingerprintStrategy(DefaultFingerprintStrategy):
     def compute_fingerprint(self, target):
       return uuid.uuid4()
 
-  def _fingerprint_strategy(self):
+  def fingerprint_strategy(self):
     if self.get_options().profile:
       # A profile is machine-specific and we assume anyone wanting a profile wants to run it here
       # and now and not accept some old result, even if on the same inputs.
@@ -585,79 +530,19 @@ def _fingerprint_strategy(self):
     else:
       return None  # Accept the default fingerprint strategy.
 
-  # Some notes on invalidation vs caching as used in `_do_run_tests` below. Here invalidation
-  # refers to executing task work in `Task.invalidated` blocks against invalid targets. Caching
-  # refers to storing the results of that work in the artifact cache using
-  # `VersionedTargetSet.results_dir`. One further bit of terminology is partition, which is the
-  # name for the set of targets passed to the `Task.invalidated` block:
-  #
-  # + Caching results for len(partition) > 1: This is trivial iff we always run all targets in
-  #   the partition, but running just invalid targets in the partition is a nicer experience (you
-  #   can whittle away at failures in a loop of `::`-style runs). Running just invalid though
-  #   requires being able to merge prior results for the partition; ie: knowing the details of
-  #   junit xml, coverage data, or using tools that do, to merge data files. The alternative is
-  #   to always run all targets in a partition if even 1 target is invalid. In this way data files
-  #   corresponding to the full partition are always generated, and so on a green partition, the
-  #   cached data files will always represent the full green run.
-  #
-  # The compromise taken here is to only cache when `all_vts == invalid_vts`; ie when the partition
-  # goes green and the run was against the full partition. A common scenario would then be:
-  #
-  # 1. Mary makes changes / adds new code and iterates `./pants test tests/python/stuff::`
-  #    gradually getting greener until finally all test targets in the `tests/python/stuff::` set
-  #    pass. She commits the green change, but there is no cached result for it since green state
-  #    for the partition was approached incrementally.
-  # 2. Jake pulls in Mary's green change and runs `./pants test tests/python/stuff::`. There is a
-  #    cache miss and he does a full local run, but since `tests/python/stuff::` is green,
-  #    `all_vts == invalid_vts` and the result is now cached for others.
-  #
-  # In this scenario, Jake will likely be a CI process, in which case human others will see a
-  # cached result from Mary's commit. It's important to note, that the CI process must run the same
-  # partition as the end user for that end user to benefit and hit the cache. This is unlikely since
-  # the only natural partitions under CI are single target ones (`--no-fast` or all targets
-  # `--fast ::`. Its unlikely an end user in a large repo will want to run `--fast ::` since `::`
-  # is probably a much wider swath of code than they're working on. As such, although `--fast`
-  # caching is supported, its unlikely to be effective. Caching is best utilized when CI and users
-  # run `--no-fast`.
-  def _do_run_tests(self, partition):
-    with self.invalidated(partition,
-                          fingerprint_strategy=self._fingerprint_strategy(),
-                          # Re-run tests when the code they test (and depend on) changes.
-                          invalidate_dependents=True) as invalidation_check:
-
-      invalid_tgts = [invalid_tgt
-                      for vts in invalidation_check.invalid_vts
-                      for invalid_tgt in vts.targets]
-
-      # Processing proceeds through:
-      # 1.) output -> workdir
-      # 2.) [iff all == invalid] workdir -> cache: We do this manually for now.
-      # 3.) [iff invalid == 0 and all > 0] cache -> workdir: Done transparently by `invalidated`.
-
-      # 1.) Write all results that will be potentially cached to workdir.
-      workdirs = _Workdirs.for_partition(self.workdir, partition)
-      result = self._run_pytest_checked(workdirs, invalid_tgts)
-
-      cache_vts = self._vts_for_partition(invalidation_check)
-      if invalidation_check.all_vts == invalidation_check.invalid_vts:
-        # 2.) The full partition was invalid, cache successful test results.
-        if result.success and self.artifact_cache_writes_enabled():
-          self.update_artifact_cache([(cache_vts, workdirs.files())])
-      elif not invalidation_check.invalid_vts:
-        # 3.) The full partition was valid, our results will have been staged for/by caching if not
-        # already local.
-        pass
-      else:
-        # The partition was partially invalid.
+  def run_tests(self, fail_fast, test_targets, workdirs):
+    try:
+      return self._run_pytest(fail_fast, test_targets, workdirs)
+    finally:
+      # Unconditionally pluck any results that an end user might need to interact with from the
+      # workdir to the locations they expect.
+      self._expose_results(test_targets, workdirs)
 
-        # We don't cache results; so others will need to re-run this partition.
-        # NB: We will presumably commit this change now though and so others will get this
-        # partition in a state that executes successfully; so when the 1st of the others
-        # executes against this partition; they will hit `all_vts == invalid_vts` and
-        # cache the results. That 1st of others is hopefully CI!
-        cache_vts.force_invalidate()
+  def result_from_error(self, error):
+    return PytestResult.from_error(error)
 
-      return result
+  def collect_files(self, workdirs):
+    return workdirs.files()
 
   def _expose_results(self, invalid_tgts, workdirs):
     external_junit_xml_dir = self.get_options().junit_xml_dir
@@ -680,20 +565,11 @@ def _expose_results(self, invalid_tgts, workdirs):
         target_dir = os.path.join(pants_distdir, 'coverage', relpath)
       mergetree(workdirs.coverage_path, target_dir)
 
-  def _run_pytest_checked(self, workdirs, targets):
-    result = self._run_pytest(workdirs, targets)
-
-    # Unconditionally pluck any results that an end user might need to interact with from the
-    # workdir to the locations they expect.
-    self._expose_results(targets, workdirs)
-
-    return result.checked()
-
-  def _run_pytest(self, workdirs, targets):
+  def _run_pytest(self, fail_fast, targets, workdirs):
     if not targets:
       return PytestResult.rc(0)
 
-    if self._run_in_chroot:
+    if self.run_tests_in_chroot:
       path_func = lambda rel_src: rel_src
     else:
       source_chroot = os.path.relpath(self._source_chroot_path(targets), get_buildroot())
@@ -721,7 +597,7 @@ def _run_pytest(self, workdirs, targets):
       # from leaking into pants test runs. See: https://github.com/pantsbuild/pants/issues/2726
       args = ['--junitxml', junitxml_path, '--confcutdir', get_buildroot(),
               '--continue-on-collection-errors']
-      if self.get_options().fail_fast:
+      if fail_fast:
         args.extend(['-x'])
       if self._debug:
         args.extend(['-s'])
@@ -778,13 +654,9 @@ def _pex_run(self, pex, workunit_name, args, env):
       process = self._spawn(pex, workunit, args, setsid=False, env=env)
       return process.wait()
 
-  @property
-  def _run_in_chroot(self):
-    return self.get_options().chroot
-
   @contextmanager
   def _maybe_run_in_chroot(self, targets):
-    if self._run_in_chroot:
+    if self.run_tests_in_chroot:
       with pushd(self._source_chroot_path(targets)):
         yield
     else:
diff --git a/src/python/pants/task/testrunner_task_mixin.py b/src/python/pants/task/testrunner_task_mixin.py
--- a/src/python/pants/task/testrunner_task_mixin.py
+++ b/src/python/pants/task/testrunner_task_mixin.py
@@ -12,6 +12,9 @@
 from threading import Timer
 
 from pants.base.exceptions import ErrorWhileTesting, TaskError
+from pants.build_graph.files import Files
+from pants.invalidation.cache_manager import VersionedTargetSet
+from pants.task.task import Task
 from pants.util.process_handler import subprocess
 
 
@@ -377,3 +380,231 @@ def _execute(self, all_targets):
     :param all_targets: list of the targets whose tests are to be run
     """
     raise NotImplementedError
+
+
+class PartitionedTestRunnerTaskMixin(TestRunnerTaskMixin, Task):
+  """A mixin for test tasks that support running tests over both individual targets and batches.
+
+  Provides support for partitioning via `--fast` (batches) and `--no-fast` (per target) options and
+  helps ensure correct caching behavior in either mode.
+
+  It's expected that mixees implement proper chrooting (see `run_tests_in_chroot`) to support
+  correct successful test result caching.
+  """
+
+  @classmethod
+  def register_options(cls, register):
+    super(PartitionedTestRunnerTaskMixin, cls).register_options(register)
+
+    # TODO(John Sirois): Implement sanity checks on options wrt caching:
+    # https://github.com/pantsbuild/pants/issues/5073
+
+    register('--fast', type=bool, default=True, fingerprint=True,
+             help='Run all tests in a single pytest invocation. If turned off, each test target '
+                  'will run in its own pytest invocation, which will be slower, but isolates '
+                  'tests from process-wide state created by tests in other targets.')
+    register('--chroot', advanced=True, fingerprint=True, type=bool, default=False,
+             help='Run tests in a chroot. Any loose files tests depend on via `{}` dependencies '
+                  'will be copied to the chroot.'
+             .format(Files.alias()))
+
+  @staticmethod
+  def _vts_for_partition(invalidation_check):
+    return VersionedTargetSet.from_versioned_targets(invalidation_check.all_vts)
+
+  def check_artifact_cache_for(self, invalidation_check):
+    # Tests generate artifacts, namely junit.xml and coverage reports, that cover the full target
+    # set whether that is all targets in the context (`--fast`) or each target individually
+    # (`--no-fast`).
+    return [self._vts_for_partition(invalidation_check)]
+
+  @property
+  def run_tests_in_chroot(self):
+    """Return `True` if tests should be run in a chroot.
+
+    Chrooted tests are expected to be run with $PWD set to a directory with only files explicitly
+    (transitively) depended on by the test targets present.
+
+    :rtype: bool
+    """
+    return self.get_options().chroot
+
+  def _execute(self, all_targets):
+    test_targets = self._get_test_targets()
+    if not test_targets:
+      return
+
+    self.context.release_lock()
+
+    per_target = not self.get_options().fast
+    fail_fast = self.get_options().fail_fast
+
+    results = {}
+    failure = False
+    with self.partitions(per_target, all_targets, test_targets) as partitions:
+      for (partition, args) in partitions():
+        try:
+          rv = self._run_partition(fail_fast, partition, *args)
+        except ErrorWhileTesting as e:
+          rv = self.result_from_error(e)
+
+        results[partition] = rv
+        if not rv.success:
+          failure = True
+          if fail_fast:
+            break
+
+      for partition in sorted(results):
+        rv = results[partition]
+        if len(partition) == 1 or rv.success:
+          log = self.context.log.info if rv.success else self.context.log.error
+          for target in partition:
+            log('{0:80}.....{1:>10}'.format(target.address.reference(), rv))
+        else:
+          # There is not much useful we can display in summary for a multi-target partition with
+          # failures without parsing those failures to link them to individual targets; ie: targets
+          # 2 and 8 failed in this partition of 10 targets.
+          # TODO(John Sirois): Punting here works for our 2 common partitionings:
+          # 1. All targets in singleton partitions
+          # 2. All targets in 1 partition
+          # PytestRun supports multiple partitions with multiple targets each when there sre
+          # multiple python source roots, and so some sort of summary for the multi-target
+          # partitions is needed: https://github.com/pantsbuild/pants/issues/5415
+          pass
+
+      msgs = [str(_rv) for _rv in results.values() if not _rv.success]
+      failed_targets = [target
+                        for _rv in results.values() if not _rv.success
+                        for target in _rv.failed_targets]
+      if len(failed_targets) > 0:
+        raise ErrorWhileTesting('\n'.join(msgs), failed_targets=failed_targets)
+      elif failure:
+        # A low-level test execution failure occurred before tests were run.
+        raise TaskError()
+
+  # Some notes on invalidation vs caching as used in `run_partition` below. Here invalidation
+  # refers to executing task work in `Task.invalidated` blocks against invalid targets. Caching
+  # refers to storing the results of that work in the artifact cache using
+  # `VersionedTargetSet.results_dir`. One further bit of terminology is partition, which is the
+  # name for the set of targets passed to the `Task.invalidated` block:
+  #
+  # + Caching results for len(partition) > 1: This is trivial iff we always run all targets in
+  #   the partition, but running just invalid targets in the partition is a nicer experience (you
+  #   can whittle away at failures in a loop of `::`-style runs). Running just invalid though
+  #   requires being able to merge prior results for the partition; ie: knowing the details of
+  #   junit xml, coverage data, or using tools that do, to merge data files. The alternative is
+  #   to always run all targets in a partition if even 1 target is invalid. In this way data files
+  #   corresponding to the full partition are always generated, and so on a green partition, the
+  #   cached data files will always represent the full green run.
+  #
+  # The compromise taken here is to only cache when `all_vts == invalid_vts`; ie when the partition
+  # goes green and the run was against the full partition. A common scenario would then be:
+  #
+  # 1. Mary makes changes / adds new code and iterates `./pants test tests/python/stuff::`
+  #    gradually getting greener until finally all test targets in the `tests/python/stuff::` set
+  #    pass. She commits the green change, but there is no cached result for it since green state
+  #    for the partition was approached incrementally.
+  # 2. Jake pulls in Mary's green change and runs `./pants test tests/python/stuff::`. There is a
+  #    cache miss and he does a full local run, but since `tests/python/stuff::` is green,
+  #    `all_vts == invalid_vts` and the result is now cached for others.
+  #
+  # In this scenario, Jake will likely be a CI process, in which case human others will see a
+  # cached result from Mary's commit. It's important to note, that the CI process must run the same
+  # partition as the end user for that end user to benefit and hit the cache. This is unlikely since
+  # the only natural partitions under CI are single target ones (`--no-fast` or all targets
+  # `--fast ::`. Its unlikely an end user in a large repo will want to run `--fast ::` since `::`
+  # is probably a much wider swath of code than they're working on. As such, although `--fast`
+  # caching is supported, its unlikely to be effective. Caching is best utilized when CI and users
+  # run `--no-fast`.
+  def _run_partition(self, fail_fast, test_targets, *args):
+    with self.invalidated(targets=test_targets,
+                          fingerprint_strategy=self.fingerprint_strategy(),
+                          # Re-run tests when the code they test (and depend on) changes.
+                          invalidate_dependents=True) as invalidation_check:
+
+      invalid_test_tgts = [invalid_test_tgt
+                           for vts in invalidation_check.invalid_vts
+                           for invalid_test_tgt in vts.targets]
+
+      # Processing proceeds through:
+      # 1.) output -> output_dir
+      # 2.) [iff all == invalid] output_dir -> cache: We do this manually for now.
+      # 3.) [iff invalid == 0 and all > 0] cache -> workdir: Done transparently by `invalidated`.
+
+      # 1.) Write all results that will be potentially cached to output_dir.
+      result = self.run_tests(fail_fast, invalid_test_tgts, *args).checked()
+
+      cache_vts = self._vts_for_partition(invalidation_check)
+      if invalidation_check.all_vts == invalidation_check.invalid_vts:
+        # 2.) All tests in the partition were invalid, cache successful test results.
+        if result.success and self.artifact_cache_writes_enabled():
+          self.update_artifact_cache([(cache_vts, self.collect_files(*args))])
+      elif not invalidation_check.invalid_vts:
+        # 3.) The full partition was valid, our results will have been staged for/by caching
+        # if not already local.
+        pass
+      else:
+        # The partition was partially invalid.
+
+        # We don't cache results; so others will need to re-run this partition.
+        # NB: We will presumably commit this change now though and so others will get this
+        # partition in a state that executes successfully; so when the 1st of the others
+        # executes against this partition; they will hit `all_vts == invalid_vts` and
+        # cache the results. That 1st of others is hopefully CI!
+        cache_vts.force_invalidate()
+
+      return result
+
+  def result_from_error(self, error):
+    """Convert an error into a test result.
+
+    :param error: The error to convert into a test result.
+    :type error: :class:`pants.base.exceptions.TaskError`
+    :returns: An unsuccessful test result.
+    :rtype: :class:`TestResult`
+    """
+    return TestResult.from_error(error)
+
+  def fingerprint_strategy(self):
+    """Return a fingerprint strategy for target fingerprinting.
+
+    :returns: A fingerprint strategy instance; by default, `None`; ie let the invalidation and
+              caching framework use the default target fingerprinter.
+    :rtype: :class:`pants.base.fingerprint_strategy.FingerprintStrategy`
+    """
+    return None
+
+  @abstractmethod
+  def partitions(self, per_target, all_targets, test_targets):
+    """Return a context manager that can be called to iterate of target partitions.
+
+    The iterator should return a 2-tuple with the partitions targets in the first slot and a tuple
+    of extra arguments needed to `run_tests` and `collect_files`.
+
+    :rtype: A context manager that is callable with no arguments; returning an iterator over
+            (partition, tuple(args))
+    """
+
+  @abstractmethod
+  def run_tests(self, fail_fast, test_targets, *args):
+    """Runs tests in the given invalid test targets.
+
+    :param bool fail_fast: `True` if the test run should fail as fast as possible.
+    :param test_targets: The test targets to run tests for.
+    :type test_targets: list of :class:`pants.build_graph.target.Target`s of the type iterated by
+                        `partitions`.
+    :param *args: Extra args associated with the partition of test targets being run as returned by
+                  the `partitions` iterator.
+    :returns: A test result summarizing the result of this test run.
+    :rtype: :class:`TestResult`
+    """
+
+  @abstractmethod
+  def collect_files(self, *args):
+    """Collects output files from a test run that should be cached.
+
+    :param *args: Extra args associated with the partition of test targets being run as returned by
+                  the `partitions` iterator.
+    :returns: A list of paths to files that should be cached.
+    :rtype: list of str
+    """


Test patch: