Add plugin hook for setting up Pytest context (pantsbuild#12091)

This new hook allows users to do things like:

- Validate that databases are running.
- Insert certain files, e.g. our `runtime_package_dependencies` feature.

In the future, we can expand it to have new capabilities like setting environment variables.

Unlike our `setup_py` kwargs plugin hook, we do allow you to run >1 plugin on the same target here. It should be safe to merge multiple plugins, and this is a useful feature.

This plugin is only implemented for Pytest, not for other test runners like `shunit2`. It seems unlikely a user would want the same setup for Pytest as other languages; and if they do, they can simply factor up a helper rule that gets used by both plugin hooks.

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/backend/python/goals/pytest_runner.py

@@ -3,6 +3,7 @@
 
 import itertools
 import logging
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import PurePath
 from typing import Optional
@@ -44,8 +45,11 @@
 )
 from pants.core.util_rules.config_files import ConfigFiles, ConfigFilesRequest
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
+from pants.engine.addresses import Address
+from pants.engine.collection import Collection
 from pants.engine.environment import CompleteEnvironment
 from pants.engine.fs import (
+    EMPTY_DIGEST,
     AddPrefix,
     CreateDigest,
     Digest,
@@ -63,21 +67,15 @@
     ProcessCacheScope,
 )
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest
-from pants.engine.unions import UnionRule
+from pants.engine.target import Target, TransitiveTargets, TransitiveTargetsRequest, WrappedTarget
+from pants.engine.unions import UnionMembership, UnionRule, union
 from pants.option.global_options import GlobalOptions
 from pants.python.python_setup import PythonSetup
 from pants.util.logging import LogLevel
 
 logger = logging.getLogger()
 
 
-# If a user wants extra pytest output (e.g., plugin output) to show up in dist/
-# they must ensure that output goes under this directory. E.g.,
-# ./pants test <target> -- --html=extra-output/report.html
-_EXTRA_OUTPUT_DIR = "extra-output"
-
-
 @dataclass(frozen=True)
 class PythonTestFieldSet(TestFieldSet):
     required_fields = (PythonTestsSources,)
@@ -96,6 +94,78 @@ def is_conftest_or_type_stub(self) -> bool:
         return file_name.name == "conftest.py" or file_name.suffix == ".pyi"
 
 
+# -----------------------------------------------------------------------------------------
+# Plugin hook
+# -----------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class PytestPluginSetup:
+    """The result of custom set up logic before Pytest runs.
+
+    Please reach out it if you would like certain functionality, such as allowing your plugin to set
+    environment variables.
+    """
+
+    digest: Digest = EMPTY_DIGEST
+
+
+@union
+@dataclass(frozen=True)  # type: ignore[misc]
+class PytestPluginSetupRequest(ABC):
+    """A request to set up the test environment before Pytest runs, e.g. to set up databases.
+
+    To use, subclass PytestPluginSetupRequest, register the rule
+    `UnionRule(PytestPluginSetupRequest, MyCustomPytestPluginSetupRequest)`, and add a rule that
+    takes your subclass as a parameter and returns `PytestPluginSetup`.
+    """
+
+    target: Target
+
+    @classmethod
+    @abstractmethod
+    def is_applicable(cls, target: Target) -> bool:
+        """Whether the setup implementation should be used for this target or not."""
+
+
+class AllPytestPluginSetups(Collection[PytestPluginSetup]):
+    pass
+
+
+# TODO: Why is this necessary? We should be able to use `PythonTestFieldSet` as the rule param.
+@dataclass(frozen=True)
+class AllPytestPluginSetupsRequest:
+    address: Address
+
+
+@rule
+async def run_all_setup_plugins(
+    request: AllPytestPluginSetupsRequest, union_membership: UnionMembership
+) -> AllPytestPluginSetups:
+    wrapped_tgt = await Get(WrappedTarget, Address, request.address)
+    applicable_setup_request_types = tuple(
+        request
+        for request in union_membership.get(PytestPluginSetupRequest)  # type: ignore[misc]
+        if request.is_applicable(wrapped_tgt.target)
+    )
+    setups = await MultiGet(
+        Get(PytestPluginSetup, PytestPluginSetupRequest, request(wrapped_tgt.target))  # type: ignore[misc]
+        for request in applicable_setup_request_types
+    )
+    return AllPytestPluginSetups(setups)
+
+
+# -----------------------------------------------------------------------------------------
+# Core logic
+# -----------------------------------------------------------------------------------------
+
+
+# If a user wants extra pytest output (e.g., plugin output) to show up in dist/
+# they must ensure that output goes under this directory. E.g.,
+# ./pants test <target> -- --html=extra-output/report.html
+_EXTRA_OUTPUT_DIR = "extra-output"
+
+
 @dataclass(frozen=True)
 class TestSetupRequest:
     field_set: PythonTestFieldSet
@@ -123,8 +193,9 @@ async def setup_pytest_for_target(
     global_options: GlobalOptions,
     complete_env: CompleteEnvironment,
 ) -> TestSetup:
-    transitive_targets = await Get(
-        TransitiveTargets, TransitiveTargetsRequest([request.field_set.address])
+    transitive_targets, plugin_setups = await MultiGet(
+        Get(TransitiveTargets, TransitiveTargetsRequest([request.field_set.address])),
+        Get(AllPytestPluginSetups, AllPytestPluginSetupsRequest(request.field_set.address)),
     )
     all_targets = transitive_targets.closure
 
@@ -151,11 +222,6 @@ async def setup_pytest_for_target(
         PythonSourceFiles, PythonSourceFilesRequest(all_targets, include_files=True)
     )
 
-    build_package_dependencies_get = Get(
-        BuiltPackageDependencies,
-        BuildPackageDependenciesRequest(request.field_set.runtime_package_dependencies),
-    )
-
     # Get the file names for the test_target so that we can specify to Pytest precisely which files
     # to test, rather than using auto-discovery.
     field_set_source_files_get = Get(SourceFiles, SourceFilesRequest([request.field_set.sources]))
@@ -165,14 +231,12 @@ async def setup_pytest_for_target(
         requirements_pex,
         prepared_sources,
         field_set_source_files,
-        built_package_dependencies,
         extra_output_directory_digest,
     ) = await MultiGet(
         pytest_pex_get,
         requirements_pex_get,
         prepared_sources_get,
         field_set_source_files_get,
-        build_package_dependencies_get,
         extra_output_directory_digest_get,
     )
 
@@ -201,7 +265,7 @@ async def setup_pytest_for_target(
                 prepared_sources.source_files.snapshot.digest,
                 config_files.snapshot.digest,
                 extra_output_directory_digest,
-                *(pkg.digest for pkg in built_package_dependencies),
+                *(plugin_setup.digest for plugin_setup in plugin_setups),
             )
         ),
     )
@@ -231,7 +295,8 @@ async def setup_pytest_for_target(
         "PYTEST_ADDOPTS": " ".join(add_opts),
         "PEX_EXTRA_SYS_PATH": ":".join(prepared_sources.source_roots),
         **test_extra_env.env,
-        # NOTE: `complete_env` intentionally after `test_extra_env` to allow overriding within `python_tests`
+        # NOTE: `complete_env` intentionally after `test_extra_env` to allow overriding within
+        # `python_tests`
         **complete_env.get_subset(request.field_set.extra_env_vars.value or ()),
     }
 
@@ -314,5 +379,31 @@ async def debug_python_test(field_set: PythonTestFieldSet) -> TestDebugRequest:
     )
 
 
+# -----------------------------------------------------------------------------------------
+# `runtime_package_dependencies` plugin
+# -----------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class RuntimePackagesPluginRequest(PytestPluginSetupRequest):
+    @classmethod
+    def is_applicable(cls, target: Target) -> bool:
+        return bool(target.get(RuntimePackageDependenciesField).value)
+
+
+@rule
+async def setup_runtime_packages(request: RuntimePackagesPluginRequest) -> PytestPluginSetup:
+    built_packages = await Get(
+        BuiltPackageDependencies,
+        BuildPackageDependenciesRequest(request.target.get(RuntimePackageDependenciesField)),
+    )
+    digest = await Get(Digest, MergeDigests(pkg.digest for pkg in built_packages))
+    return PytestPluginSetup(digest)
+
+
 def rules():
-    return [*collect_rules(), UnionRule(TestFieldSet, PythonTestFieldSet)]
+    return [
+        *collect_rules(),
+        UnionRule(TestFieldSet, PythonTestFieldSet),
+        UnionRule(PytestPluginSetupRequest, RuntimePackagesPluginRequest),
+    ]

src/python/pants/backend/python/goals/pytest_runner_integration_test.py

@@ -13,7 +13,11 @@
 from pants.backend.python.dependency_inference import rules as dependency_inference_rules
 from pants.backend.python.goals import package_pex_binary, pytest_runner
 from pants.backend.python.goals.coverage_py import create_or_update_coverage_config
-from pants.backend.python.goals.pytest_runner import PythonTestFieldSet
+from pants.backend.python.goals.pytest_runner import (
+    PytestPluginSetup,
+    PytestPluginSetupRequest,
+    PythonTestFieldSet,
+)
 from pants.backend.python.target_types import (
     PexBinary,
     PythonLibrary,
@@ -29,9 +33,11 @@
 )
 from pants.core.util_rules import config_files, distdir
 from pants.engine.addresses import Address
-from pants.engine.fs import DigestContents
+from pants.engine.fs import CreateDigest, Digest, DigestContents, FileContent
 from pants.engine.process import InteractiveRunner
+from pants.engine.rules import Get, rule
 from pants.engine.target import Target
+from pants.engine.unions import UnionRule
 from pants.testutil.python_interpreter_selection import skip_unless_python27_and_python3_present
 from pants.testutil.rule_runner import QueryRule, RuleRunner, mock_console
 
@@ -420,7 +426,43 @@ def test_args():
     assert result.exit_code == 0
 
 
-def test_runtime_package_dependency(rule_runner: RuleRunner) -> None:
+class UsedPlugin(PytestPluginSetupRequest):
+    @classmethod
+    def is_applicable(cls, target: Target) -> bool:
+        return True
+
+
+class UnusedPlugin(PytestPluginSetupRequest):
+    @classmethod
+    def is_applicable(cls, target: Target) -> bool:
+        return False
+
+
+@rule
+async def used_plugin(_: UsedPlugin) -> PytestPluginSetup:
+    digest = await Get(Digest, CreateDigest([FileContent("used.txt", b"")]))
+    return PytestPluginSetup(digest=digest)
+
+
+@rule
+async def unused_plugin(_: UnusedPlugin) -> PytestPluginSetup:
+    digest = await Get(Digest, CreateDigest([FileContent("unused.txt", b"")]))
+    return PytestPluginSetup(digest=digest)
+
+
+def test_setup_plugins_and_runtime_package_dependency(rule_runner: RuleRunner) -> None:
+    # We test both the generic `PytestPluginSetup` mechanism and our `runtime_package_dependencies`
+    # feature in the same test to confirm multiple plugins can be used on the same target.
+    rule_runner = RuleRunner(
+        rules=[
+            *rule_runner.rules,
+            used_plugin,
+            unused_plugin,
+            UnionRule(PytestPluginSetupRequest, UsedPlugin),
+            UnionRule(PytestPluginSetupRequest, UnusedPlugin),
+        ],
+        target_types=rule_runner.target_types,
+    )
     rule_runner.write_files(
         {
             f"{PACKAGE}/say_hello.py": "print('Hello, test!')",
@@ -430,12 +472,17 @@ def test_runtime_package_dependency(rule_runner: RuleRunner) -> None:
                 import subprocess
 
                 def test_embedded_binary():
+                    assert os.path.exists("bin.pex")
                     assert b"Hello, test!" in subprocess.check_output(args=['./bin.pex'])
 
                     # Ensure that we didn't accidentally pull in the binary's sources. This is a
                     # special type of dependency that should not be included with the rest of the
                     # normal dependencies.
-                    assert os.path.exists("{PACKAGE}/say_hello.py") is False
+                    assert not os.path.exists("{PACKAGE}/say_hello.py")
+
+                def test_additional_plugins():
+                    assert os.path.exists("used.txt")
+                    assert not os.path.exists("unused.txt")
                 """
             ),
             f"{PACKAGE}/BUILD": dedent(
@@ -448,7 +495,7 @@ def test_embedded_binary():
         }
     )
     tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="test_binary_call.py"))
-    result = run_pytest(rule_runner, tgt, extra_args=["--pytest-args='-s'"])
+    result = run_pytest(rule_runner, tgt)
     assert result.exit_code == 0
 
 

src/python/pants/testutil/rule_runner.py

@@ -32,7 +32,7 @@
 from pants.engine.rules import QueryRule as QueryRule
 from pants.engine.rules import Rule
 from pants.engine.target import Target, WrappedTarget
-from pants.engine.unions import UnionMembership
+from pants.engine.unions import UnionMembership, UnionRule
 from pants.init.engine_initializer import EngineInitializer
 from pants.init.logging import initialize_stdio, stdio_destination
 from pants.option.global_options import (
@@ -191,8 +191,8 @@ def pants_workdir(self) -> str:
         return os.path.join(self.build_root, ".pants.d")
 
     @property
-    def rules(self) -> FrozenOrderedSet[Rule]:
-        return self.build_config.rules
+    def rules(self) -> FrozenOrderedSet[Rule | UnionRule]:
+        return FrozenOrderedSet([*self.build_config.rules, *self.build_config.union_rules])
 
     @property
     def target_types(self) -> FrozenOrderedSet[Type[Target]]: