# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import annotations

import itertools
import json
import logging
import os
import shlex
from abc import ABC, ABCMeta
from collections.abc import Coroutine, Iterable, Sequence
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import PurePath
from typing import Any, ClassVar, TypeVar, cast

from pants.core.environments.rules import (
    ChosenLocalEnvironmentName,
    EnvironmentName,
    SingleEnvironmentNameRequest,
    resolve_single_environment_name,
)
from pants.core.goals.multi_tool_goal_helper import SkippableSubsystem
from pants.core.goals.package import (
    BuiltPackage,
    EnvironmentAwarePackageRequest,
    PackageFieldSet,
    environment_aware_package,
)
from pants.core.subsystems.debug_adapter import DebugAdapterSubsystem
from pants.core.util_rules.distdir import DistDir
from pants.core.util_rules.env_vars import environment_vars_subset
from pants.core.util_rules.partitions import (
    PartitionerType,
    PartitionMetadataT,
    Partitions,
    _BatchBase,
    _PartitionFieldSetsRequestBase,
)
from pants.engine.addresses import Address
from pants.engine.collection import Collection
from pants.engine.console import Console
from pants.engine.desktop import OpenFilesRequest, find_open_program
from pants.engine.engine_aware import EngineAwareReturnType
from pants.engine.env_vars import EXTRA_ENV_VARS_USAGE_HELP, EnvironmentVars, EnvironmentVarsRequest
from pants.engine.fs import EMPTY_FILE_DIGEST, FileDigest, MergeDigests, Snapshot, Workspace
from pants.engine.goal import Goal, GoalSubsystem
from pants.engine.internals.graph import find_valid_field_sets, resolve_targets
from pants.engine.internals.session import RunId
from pants.engine.internals.specs_rules import find_valid_field_sets_for_target_roots
from pants.engine.intrinsics import merge_digests, run_interactive_process_in_environment
from pants.engine.process import FallibleProcessResult, InteractiveProcess, ProcessResultMetadata
from pants.engine.rules import collect_rules, concurrently, goal_rule, implicitly, rule
from pants.engine.target import (
    FieldSet,
    FieldSetsPerTargetRequest,
    IntField,
    NoApplicableTargetsBehavior,
    SourcesField,
    SpecialCasedDependencies,
    StringField,
    StringSequenceField,
    TargetRootsToFieldSets,
    TargetRootsToFieldSetsRequest,
    ValidNumbers,
    parse_shard_spec,
)
from pants.engine.unions import UnionMembership, UnionRule, distinct_union_type_per_subclass, union
from pants.option.option_types import BoolOption, EnumOption, IntOption, StrListOption, StrOption
from pants.util.collections import partition_sequentially
from pants.util.dirutil import safe_open
from pants.util.docutil import bin_name
from pants.util.logging import LogLevel
from pants.util.memo import memoized, memoized_property
from pants.util.meta import classproperty
from pants.util.strutil import Simplifier, help_text, softwrap

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class TestResult(EngineAwareReturnType):
    # A None exit_code indicates a backend that performs its own test discovery/selection
    # (rather than delegating that to the underlying test tool), and discovered no tests.
    exit_code: int | None
    stdout_bytes: bytes
    stdout_digest: FileDigest
    stderr_bytes: bytes
    stderr_digest: FileDigest
    addresses: tuple[Address, ...]
    output_setting: ShowOutput
    # A None result_metadata indicates a backend that performs its own test discovery/selection
    # and either discovered no tests, or encountered an error, such as a compilation error, in
    # the attempt.
    result_metadata: ProcessResultMetadata | None  # TODO: Merge elapsed MS of all subproceses
    partition_description: str | None = None

    coverage_data: CoverageData | None = None
    # TODO: Rename this to `reports`. There is no guarantee that every language will produce
    #  XML reports, or only XML reports.
    xml_results: Snapshot | None = None
    # Any extra output (such as from plugins) that the test runner was configured to output.
    extra_output: Snapshot | None = None
    # True if the core test rules should log that extra output was written.
    log_extra_output: bool = False
    # All results including failed attempts
    process_results: tuple[FallibleProcessResult, ...] = field(default_factory=tuple)

    output_simplifier: Simplifier = Simplifier()

    # Prevent this class from being detected by pytest as a test class.
    __test__ = False

    @staticmethod
    def no_tests_found(address: Address, output_setting: ShowOutput) -> TestResult:
        """Used when we do test discovery ourselves, and we didn't find any."""
        return TestResult(
            exit_code=None,
            stdout_bytes=b"",
            stderr_bytes=b"",
            stdout_digest=EMPTY_FILE_DIGEST,
            stderr_digest=EMPTY_FILE_DIGEST,
            addresses=(address,),
            output_setting=output_setting,
            result_metadata=None,
        )

    @staticmethod
    def no_tests_found_in_batch(
        batch: TestRequest.Batch[_TestFieldSetT, Any], output_setting: ShowOutput
    ) -> TestResult:
        """Used when we do test discovery ourselves, and we didn't find any."""
        return TestResult(
            exit_code=None,
            stdout_bytes=b"",
            stderr_bytes=b"",
            stdout_digest=EMPTY_FILE_DIGEST,
            stderr_digest=EMPTY_FILE_DIGEST,
            addresses=tuple(field_set.address for field_set in batch.elements),
            output_setting=output_setting,
            result_metadata=None,
            partition_description=batch.partition_metadata.description,
        )

    @staticmethod
    def from_fallible_process_result(
        process_results: tuple[FallibleProcessResult, ...],
        address: Address,
        output_setting: ShowOutput,
        *,
        coverage_data: CoverageData | None = None,
        xml_results: Snapshot | None = None,
        extra_output: Snapshot | None = None,
        log_extra_output: bool = False,
        output_simplifier: Simplifier = Simplifier(),
    ) -> TestResult:
        process_result = process_results[-1]
        return TestResult(
            exit_code=process_result.exit_code,
            stdout_bytes=process_result.stdout,
            stdout_digest=process_result.stdout_digest,
            stderr_bytes=process_result.stderr,
            stderr_digest=process_result.stderr_digest,
            addresses=(address,),
            output_setting=output_setting,
            result_metadata=process_result.metadata,
            coverage_data=coverage_data,
            xml_results=xml_results,
            extra_output=extra_output,
            log_extra_output=log_extra_output,
            process_results=process_results,
            output_simplifier=output_simplifier,
        )

    @staticmethod
    def from_batched_fallible_process_result(
        process_results: tuple[FallibleProcessResult, ...],
        batch: TestRequest.Batch[_TestFieldSetT, Any],
        output_setting: ShowOutput,
        *,
        coverage_data: CoverageData | None = None,
        xml_results: Snapshot | None = None,
        extra_output: Snapshot | None = None,
        log_extra_output: bool = False,
        output_simplifier: Simplifier = Simplifier(),
    ) -> TestResult:
        process_result = process_results[-1]
        return TestResult(
            exit_code=process_result.exit_code,
            stdout_bytes=process_result.stdout,
            stdout_digest=process_result.stdout_digest,
            stderr_bytes=process_result.stderr,
            stderr_digest=process_result.stderr_digest,
            addresses=tuple(field_set.address for field_set in batch.elements),
            output_setting=output_setting,
            result_metadata=process_result.metadata,
            coverage_data=coverage_data,
            xml_results=xml_results,
            extra_output=extra_output,
            log_extra_output=log_extra_output,
            output_simplifier=output_simplifier,
            partition_description=batch.partition_metadata.description,
            process_results=process_results,
        )

    @property
    def description(self) -> str:
        if len(self.addresses) == 1:
            return self.addresses[0].spec

        return f"{self.addresses[0].spec} and {len(self.addresses) - 1} other files"

    @property
    def path_safe_description(self) -> str:
        if len(self.addresses) == 1:
            return self.addresses[0].path_safe_spec

        return f"{self.addresses[0].path_safe_spec}+{len(self.addresses) - 1}"

    def __lt__(self, other: Any) -> bool:
        """We sort first by exit code, then alphanumerically within each group."""
        if not isinstance(other, TestResult):
            return NotImplemented
        if self.exit_code == other.exit_code:
            return self.description < other.description
        if self.exit_code is None:
            return True
        if other.exit_code is None:
            return False
        return abs(self.exit_code) < abs(other.exit_code)

    def artifacts(self) -> dict[str, FileDigest | Snapshot] | None:
        output: dict[str, FileDigest | Snapshot] = {
            "stdout": self.stdout_digest,
            "stderr": self.stderr_digest,
        }
        if self.xml_results:
            output["xml_results"] = self.xml_results
        return output

    def level(self) -> LogLevel:
        if self.exit_code is None:
            return LogLevel.DEBUG
        return LogLevel.INFO if self.exit_code == 0 else LogLevel.ERROR

    def _simplified_output(self, v: bytes) -> str:
        return self.output_simplifier.simplify(v.decode(errors="replace"))

    @memoized_property
    def stdout_simplified_str(self) -> str:
        return self._simplified_output(self.stdout_bytes)

    @memoized_property
    def stderr_simplified_str(self) -> str:
        return self._simplified_output(self.stderr_bytes)

    def message(self) -> str:
        if self.exit_code is None:
            return "no tests found."
        status = "succeeded" if self.exit_code == 0 else f"failed (exit code {self.exit_code})"
        message = f"{status}."
        if self.partition_description:
            message += f"\nPartition: {self.partition_description}"
        if self.output_setting == ShowOutput.NONE or (
            self.output_setting == ShowOutput.FAILED and self.exit_code == 0
        ):
            return message
        output = ""
        if self.stdout_bytes:
            output += f"\n{self.stdout_simplified_str}"
        if self.stderr_bytes:
            output += f"\n{self.stderr_simplified_str}"
        if output:
            output = f"{output.rstrip()}\n\n"
        return f"{message}{output}"

    def metadata(self) -> dict[str, Any]:
        return {"addresses": [address.spec for address in self.addresses]}

    def cacheable(self) -> bool:
        """Is marked uncacheable to ensure that it always renders."""
        return False


class ShowOutput(Enum):
    """Which tests to emit detailed output for."""

    ALL = "all"
    FAILED = "failed"
    NONE = "none"


@dataclass(frozen=True)
class TestDebugRequest:
    process: InteractiveProcess

    # Prevent this class from being detected by pytest as a test class.
    __test__ = False


class TestDebugAdapterRequest(TestDebugRequest):
    """Like TestDebugRequest, but launches the test process using the relevant Debug Adapter server.

    The process should be launched waiting for the client to connect.
    """


@union
@dataclass(frozen=True)
class TestFieldSet(FieldSet, metaclass=ABCMeta):
    """The fields necessary to run tests on a target."""

    sources: SourcesField

    __test__ = False


_TestFieldSetT = TypeVar("_TestFieldSetT", bound=TestFieldSet)


@union
class TestRequest:
    """Base class for plugin types wanting to be run as part of `test`.

    Plugins should define a new type which subclasses this type, and set the
    appropriate class variables.
    E.g.
        class DryCleaningRequest(TestRequest):
            tool_subsystem = DryCleaningSubsystem
            field_set_type = DryCleaningFieldSet

    Then register the rules which tell Pants about your plugin.
    E.g.
        def rules():
            return [
                *collect_rules(),
                *DryCleaningRequest.rules(),
            ]
    """

    tool_subsystem: ClassVar[type[SkippableSubsystem]]
    field_set_type: ClassVar[type[TestFieldSet]]
    partitioner_type: ClassVar[PartitionerType] = PartitionerType.DEFAULT_ONE_PARTITION_PER_INPUT

    supports_debug: ClassVar[bool] = False
    supports_debug_adapter: ClassVar[bool] = False

    __test__ = False

    @classproperty
    def tool_name(cls) -> str:
        return cls.tool_subsystem.options_scope

    @distinct_union_type_per_subclass(in_scope_types=[EnvironmentName])
    class PartitionRequest(_PartitionFieldSetsRequestBase[_TestFieldSetT]):
        def metadata(self) -> dict[str, Any]:
            return {"addresses": [field_set.address.spec for field_set in self.field_sets]}

    @distinct_union_type_per_subclass(in_scope_types=[EnvironmentName])
    class Batch(_BatchBase[_TestFieldSetT, PartitionMetadataT]):
        @property
        def single_element(self) -> _TestFieldSetT:
            """Return the single element of this batch.

            NOTE: Accessing this property will raise a `TypeError` if this `Batch` contains
            >1 elements. It is only safe to be used by test runners utilizing the "default"
            one-input-per-partition partitioner type.
            """

            if len(self.elements) != 1:
                description = ""
                if self.partition_metadata.description:
                    description = f" from partition '{self.partition_metadata.description}'"
                raise TypeError(
                    f"Expected a single element in batch{description}, but found {len(self.elements)}"
                )

            return self.elements[0]

        @property
        def description(self) -> str:
            if self.partition_metadata and self.partition_metadata.description:
                return f"test batch from partition '{self.partition_metadata.description}'"
            return "test batch"

        def debug_hint(self) -> str:
            if len(self.elements) == 1:
                return self.elements[0].address.spec

            return f"{self.elements[0].address.spec} and {len(self.elements) - 1} other files"

        def metadata(self) -> dict[str, Any]:
            return {
                "addresses": [field_set.address.spec for field_set in self.elements],
                "partition_description": self.partition_metadata.description,
            }

    @classmethod
    def rules(cls) -> Iterable:
        yield from cls.partitioner_type.default_rules(cls, by_file=False)

        yield UnionRule(TestFieldSet, cls.field_set_type)
        yield UnionRule(TestRequest, cls)
        yield UnionRule(TestRequest.PartitionRequest, cls.PartitionRequest)
        yield UnionRule(TestRequest.Batch, cls.Batch)

        if not cls.supports_debug:
            yield from _unsupported_debug_rules(cls)

        if not cls.supports_debug_adapter:
            yield from _unsupported_debug_adapter_rules(cls)


@rule(polymorphic=True)
async def partition_tests(req: TestRequest.PartitionRequest) -> Partitions:
    raise NotImplementedError()


@rule(polymorphic=True)
async def test_batch_to_debug_request(batch: TestRequest.Batch) -> TestDebugRequest:
    raise NotImplementedError()


@rule(polymorphic=True)
async def test_batch_to_debug_adapter_request(batch: TestRequest.Batch) -> TestDebugAdapterRequest:
    raise NotImplementedError()


@rule(polymorphic=True)
async def run_test_batch(batch: TestRequest.Batch) -> TestResult:
    raise NotImplementedError()


class CoverageData(ABC):
    """Base class for inputs to a coverage report.

    Subclasses should add whichever fields they require - snapshots of coverage output, XML files,
    etc.
    """


_CD = TypeVar("_CD", bound=CoverageData)


@union(in_scope_types=[EnvironmentName])
class CoverageDataCollection(Collection[_CD]):
    element_type: ClassVar[type[_CD]]


@dataclass(frozen=True)
class CoverageReport(ABC):
    """Represents a code coverage report that can be materialized to the terminal or disk."""

    # Some coverage systems can determine, based on a configurable threshold, whether coverage
    # was sufficient or not. The test goal will fail the build if coverage was deemed insufficient.
    coverage_insufficient: bool

    def materialize(self, console: Console, workspace: Workspace) -> PurePath | None:
        """Materialize this code coverage report to the terminal or disk.

        :param console: A handle to the terminal.
        :param workspace: A handle to local disk.
        :return: If a report was materialized to disk, the path of the file in the report one might
                 open first to start examining the report.
        """
        ...

    def get_artifact(self) -> tuple[str, Snapshot] | None:
        return None


@dataclass(frozen=True)
class ConsoleCoverageReport(CoverageReport):
    """Materializes a code coverage report to the terminal."""

    report: str

    def materialize(self, console: Console, workspace: Workspace) -> None:
        console.print_stderr(f"\n{self.report}")
        return None


@dataclass(frozen=True)
class FilesystemCoverageReport(CoverageReport):
    """Materializes a code coverage report to disk."""

    result_snapshot: Snapshot
    directory_to_materialize_to: PurePath
    report_file: PurePath | None
    report_type: str

    def materialize(self, console: Console, workspace: Workspace) -> PurePath | None:
        workspace.write_digest(
            self.result_snapshot.digest, path_prefix=str(self.directory_to_materialize_to)
        )
        console.print_stderr(
            f"\nWrote {self.report_type} coverage report to `{self.directory_to_materialize_to}`"
        )
        return self.report_file

    def get_artifact(self) -> tuple[str, Snapshot] | None:
        return f"coverage_{self.report_type}", self.result_snapshot


@dataclass(frozen=True)
class CoverageReports(EngineAwareReturnType):
    reports: tuple[CoverageReport, ...]

    @property
    def coverage_insufficient(self) -> bool:
        """Whether to fail the build due to insufficient coverage."""
        return any(report.coverage_insufficient for report in self.reports)

    def materialize(self, console: Console, workspace: Workspace) -> tuple[PurePath, ...]:
        report_paths = []
        for report in self.reports:
            report_path = report.materialize(console, workspace)
            if report_path:
                report_paths.append(report_path)
        return tuple(report_paths)

    def artifacts(self) -> dict[str, Snapshot | FileDigest] | None:
        artifacts: dict[str, Snapshot | FileDigest] = {}
        for report in self.reports:
            artifact = report.get_artifact()
            if not artifact:
                continue
            artifacts[artifact[0]] = artifact[1]
        return artifacts or None


@rule(polymorphic=True)
async def create_coverage_report(req: CoverageDataCollection) -> CoverageReports:
    raise NotImplementedError()


class TestSubsystem(GoalSubsystem):
    name = "test"
    help = "Run tests."

    # Prevent this class from being detected by pytest as a test class.
    __test__ = False

    @classmethod
    def activated(cls, union_membership: UnionMembership) -> bool:
        return TestRequest in union_membership

    class EnvironmentAware:
        extra_env_vars = StrListOption(
            help=softwrap(
                f"""
                Additional environment variables to include in test processes.

                {EXTRA_ENV_VARS_USAGE_HELP}
                """
            ),
        )

    debug = BoolOption(
        default=False,
        help=softwrap(
            """
            Run tests sequentially in an interactive process. This is necessary, for
            example, when you add breakpoints to your code.
            """
        ),
    )
    # See also `run.py`'s same option
    debug_adapter = BoolOption(
        default=False,
        help=softwrap(
            """
            Run tests sequentially in an interactive process, using a Debug Adapter
            (https://microsoft.github.io/debug-adapter-protocol/) for the language if supported.

            The interactive process used will be immediately blocked waiting for a client before
            continuing.

            This option implies `--debug`.
            """
        ),
    )
    force = BoolOption(
        default=False,
        help="Force the tests to run, even if they could be satisfied from cache.",
    )
    output = EnumOption(
        default=ShowOutput.FAILED,
        help="Show stdout/stderr for these tests.",
    )
    use_coverage = BoolOption(
        default=False,
        help="Generate a coverage report if the test runner supports it.",
    )
    open_coverage = BoolOption(
        default=False,
        help=softwrap(
            """
            If a coverage report file is generated, open it on the local system if the
            system supports this.
            """
        ),
    )
    report = BoolOption(default=False, advanced=True, help="Write test reports to `--report-dir`.")
    default_report_path = str(PurePath("{distdir}", "test", "reports"))
    _report_dir = StrOption(
        default=default_report_path,
        advanced=True,
        help="Path to write test reports to. Must be relative to the build root.",
    )
    shard = StrOption(
        default="",
        help=softwrap(
            """
            A shard specification of the form "k/N", where N is a positive integer and k is a
            non-negative integer less than N.

            If set, the request input targets will be deterministically partitioned into N disjoint
            subsets of roughly equal size, and only the k'th subset will be used, with all others
            discarded.

            Useful for splitting large numbers of test files across multiple machines in CI.
            For example, you can run three shards with `--shard=0/3`, `--shard=1/3`, `--shard=2/3`.

            Note that the shards are roughly equal in size as measured by number of files.
            No attempt is made to consider the size of different files, the time they have
            taken to run in the past, or other such sophisticated measures.
            """
        ),
    )
    timeouts = BoolOption(
        default=True,
        help=softwrap(
            """
            Enable test target timeouts. If timeouts are enabled then test targets with a
            `timeout=` parameter set on their target will time out after the given number of
            seconds if not completed. If no timeout is set, then either the default timeout
            is used or no timeout is configured.
            """
        ),
    )
    timeout_default = IntOption(
        default=None,
        advanced=True,
        help=softwrap(
            """
            The default timeout (in seconds) for a test target if the `timeout` field is not
            set on the target.
            """
        ),
    )
    timeout_maximum = IntOption(
        default=None,
        advanced=True,
        help="The maximum timeout (in seconds) that may be used on a test target.",
    )
    _attempts_default = IntOption(
        default=1,
        help=softwrap(
            """
            The number of attempts to run tests, in case of a test failure.
            Tests that were retried will include the number of attempts in the summary output.
            """
        ),
    )

    batch_size = IntOption(
        "--batch-size",
        default=128,
        advanced=True,
        help=softwrap(
            """
            The target maximum number of files to be included in each run of batch-enabled
            test runners.

            Some test runners can execute tests from multiple files in a single run. Test
            implementations will return all tests that _can_ run together as a single group -
            and then this may be further divided into smaller batches, based on this option.
            This is done:

              1. to avoid OS argument length limits (in processes which don't support argument files)
              2. to support more stable cache keys than would be possible if all files were operated \
                 on in a single batch
              3. to allow for parallelism in test runners which don't have internal \
                 parallelism, or -- if they do support internal parallelism -- to improve scheduling \
                 behavior when multiple processes are competing for cores and so internal parallelism \
                 cannot be used perfectly

            In order to improve cache hit rates (see 2.), batches are created at stable boundaries,
            and so this value is only a "target" max batch size (rather than an exact value).

            NOTE: This parameter has no effect on test runners/plugins that do not implement support
            for batched testing.
            """
        ),
    )

    show_rerun_command = BoolOption(
        default="CI" in os.environ,
        advanced=True,
        help=softwrap(
            f"""
            If tests fail, show an appropriate `{bin_name()} {name} ...` invocation to rerun just
            those tests.

            This is to make it easy to run those tests on a new machine (for instance, run tests
            locally if they fail in CI): caching of successful tests means that rerunning the exact
            same command on the same machine will already automatically only rerun the failures.

            This defaults to `True` when running in CI (as determined by the `CI` environment
            variable being set) but `False` elsewhere.
            """
        ),
    )
    experimental_report_test_result_info = BoolOption(
        default=False,
        advanced=True,
        help=softwrap(
            """
            Report information about the test results.

            For now, it reports only the source from where the test results were fetched. When running tests,
            they may be executed locally or remotely, but if there are results of previous runs available,
            they may be retrieved from the local or remote cache, or be memoized. Knowing where the test
            results come from might be useful when evaluating the efficiency of the cache and the nature of
            the changes in the source code that may lead to frequent cache invalidations.
            """
        ),
    )

    def report_dir(self, distdir: DistDir) -> PurePath:
        return PurePath(self._report_dir.format(distdir=distdir.relpath))

    @property
    def attempts_default(self):
        if self._attempts_default < 1:
            raise ValueError(
                "The `--test-attempts-default` option must have a value equal or greater than 1. "
                f"Instead, it was set to {self._attempts_default}."
            )
        return self._attempts_default


class Test(Goal):
    __test__ = False

    subsystem_cls = TestSubsystem
    environment_behavior = Goal.EnvironmentBehavior.USES_ENVIRONMENTS


class TestTimeoutField(IntField, metaclass=ABCMeta):
    """Base field class for implementing timeouts for test targets.

    Each test target that wants to implement a timeout needs to provide with its own concrete field
    class extending this one.
    """

    __test__ = False

    alias = "timeout"
    required = False
    valid_numbers = ValidNumbers.positive_only
    help = help_text(
        """
        A timeout (in seconds) used by each test file belonging to this target.

        If unset, will default to `[test].timeout_default`; if that option is also unset,
        then the test will never time out. Will never exceed `[test].timeout_maximum`. Only
        applies if the option `--test-timeouts` is set to true (the default).
        """
    )

    def calculate_from_global_options(self, test: TestSubsystem) -> int | None:
        if not test.timeouts:
            return None
        if self.value is None:
            if test.timeout_default is None:
                return None
            result = test.timeout_default
        else:
            result = self.value
        if test.timeout_maximum is not None:
            return min(result, test.timeout_maximum)
        return result


class TestExtraEnvVarsField(StringSequenceField, metaclass=ABCMeta):
    alias = "extra_env_vars"
    help = help_text(
        f"""
        Additional environment variables to include in test processes.

        {EXTRA_ENV_VARS_USAGE_HELP}

        This will be merged with and override values from `[test].extra_env_vars`.
        """
    )

    def sorted(self) -> tuple[str, ...]:
        return tuple(sorted(self.value or ()))


class TestsBatchCompatibilityTagField(StringField, metaclass=ABCMeta):
    alias = "batch_compatibility_tag"

    @classmethod
    def format_help(cls, target_name: str, test_runner_name: str) -> str:
        return f"""
        An arbitrary value used to mark the test files belonging to this target as valid for
        batched execution.

        It's _sometimes_ safe to run multiple `{target_name}`s within a single test runner process,
        and doing so can give significant wins by allowing reuse of expensive test setup /
        teardown logic. To opt into this behavior, set this field to an arbitrary non-empty
        string on all the `{target_name}` targets that are safe/compatible to run in the same
        process.

        If this field is left unset on a target, the target is assumed to be incompatible with
        all others and will run in a dedicated `{test_runner_name}` process.

        If this field is set on a target, and its value is different from the value on some
        other test `{target_name}`, then the two targets are explicitly incompatible and are guaranteed
        to not run in the same `{test_runner_name}` process.

        If this field is set on a target, and its value is the same as the value on some other
        `{target_name}`, then the two targets are explicitly compatible and _may_ run in the same
        test runner process. Compatible tests may not end up in the same test runner batch if:

          * There are "too many" compatible tests in a partition, as determined by the \
            `[test].batch_size` config parameter, or
          * Compatible tests have some incompatibility in Pants metadata (i.e. different \
            `resolve`s or `extra_env_vars`).

        When tests with the same `batch_compatibility_tag` have incompatibilities in some other
        Pants metadata, they will be automatically split into separate batches. This way you can
        set a high-level `batch_compatibility_tag` using `__defaults__` and then have tests
        continue to work as you tweak BUILD metadata on specific targets.
        """


async def _get_test_batches(
    core_request_types: Iterable[type[TestRequest]],
    targets_to_field_sets: TargetRootsToFieldSets,
    local_environment_name: ChosenLocalEnvironmentName,
    test_subsystem: TestSubsystem,
) -> list[TestRequest.Batch]:
    def partitions_call(request_type: type[TestRequest]) -> Coroutine[Any, Any, Partitions]:
        partition_type = cast(TestRequest, request_type)
        field_set_type = partition_type.field_set_type
        applicable_field_sets: list[TestFieldSet] = []
        for target, field_sets in targets_to_field_sets.mapping.items():
            if field_set_type.is_applicable(target):
                applicable_field_sets.extend(field_sets)

        partition_request = partition_type.PartitionRequest(tuple(applicable_field_sets))
        return partition_tests(
            **implicitly(
                {
                    partition_request: TestRequest.PartitionRequest,
                    local_environment_name.val: EnvironmentName,
                },
            )
        )

    all_partitions = await concurrently(
        partitions_call(request_type) for request_type in core_request_types
    )

    return [
        request_type.Batch(
            cast(TestRequest, request_type).tool_name, tuple(batch), partition.metadata
        )
        for request_type, partitions in zip(core_request_types, all_partitions)
