Distribution Scheduling

class Scheduling(Protocol):
    """Protocol for all test distribution schedulers."""
    
    @property
    def nodes(self) -> list[WorkerController]:
        """List of active worker nodes."""
    
    @property
    def collection_is_completed(self) -> bool:
        """True if test collection is complete on all nodes."""
    
    @property
    def tests_finished(self) -> bool:
        """True if all tests have been executed."""
    
    @property
    def has_pending(self) -> bool:
        """True if there are pending tests to be scheduled."""
    
    def add_node(self, node: WorkerController) -> None:
        """Add a worker node to the scheduler."""
    
    def add_node_collection(
        self,
        node: WorkerController,
        collection: Sequence[str],
    ) -> None:
        """Add collected test items from a worker node."""
    
    def mark_test_complete(
        self,
        node: WorkerController,
        item_index: int,
        duration: float = 0,
    ) -> None:
        """Mark a test as completed by a worker."""
    
    def mark_test_pending(self, item: str) -> None:
        """Mark a test as pending (needs to be rescheduled)."""
    
    def remove_pending_tests_from_node(
        self,
        node: WorkerController,
        indices: Sequence[int],
    ) -> None:
        """Remove pending tests from a specific worker node."""
    
    def remove_node(self, node: WorkerController) -> str | None:
        """Remove a worker node from the scheduler."""
    
    def schedule(self) -> None:
        """Execute the scheduling algorithm to distribute tests."""

class LoadScheduling:
    """
    Load balance by sending any pending test to any available worker.
    
    Tests are distributed one at a time to the next available worker,
    providing good load balancing for test suites with varying execution times.
    """

class EachScheduling:
    """
    Send each test to all available workers.
    
    Every test is executed on every worker, useful for testing
    across different environments or configurations.
    """

class LoadScopeScheduling:
    """
    Load balance by sending pending groups of tests in the same scope
    to any available worker.
    
    Groups tests by scope (typically class or module) to improve
    fixture reuse and reduce setup/teardown overhead.
    """

class LoadFileScheduling:
    """
    Load balance by sending tests grouped by file to any available worker.
    
    All tests from the same file are sent to the same worker,
    optimizing for file-level fixtures and reducing I/O overhead.
    """

class LoadGroupScheduling:
    """
    Like LoadScheduling, but sends tests marked with 'xdist_group' 
    to the same worker.
    
    Tests marked with the same xdist_group name will always run
    on the same worker, useful for tests that share state or resources.
    """

class WorkStealingScheduling:
    """
    Split the test suite evenly between available workers initially,
    then rebalance when any worker runs out of tests.
    
    Provides good initial distribution with dynamic rebalancing
    for optimal resource utilization.
    """

def pytest_xdist_make_scheduler(
    config: pytest.Config, log: Producer
) -> Scheduling | None:
    """
    Hook for creating custom scheduler implementations.
    
    Args:
        config: pytest configuration object
        log: logging producer for scheduler messages
        
    Returns:
        Custom scheduler instance or None to use default
    """

# Load balancing (default with -n)
pytest -n 4

# Explicit load balancing
pytest -n 4 --dist load

# Each test on all workers
pytest -n 4 --dist each

# Group by scope
pytest -n 4 --dist loadscope

# Group by file
pytest -n 4 --dist loadfile

# Respect xdist_group markers
pytest -n 4 --dist loadgroup

# Work stealing
pytest -n 4 --dist worksteal

# Control loadscope reordering
pytest -n 4 --dist loadscope --loadscope-reorder
pytest -n 4 --dist loadscope --no-loadscope-reorder

import pytest

# Tests with same group run on same worker
@pytest.mark.xdist_group(name="database")
def test_db_setup():
    # Database initialization
    pass

@pytest.mark.xdist_group(name="database")  
def test_db_query():
    # Can rely on setup from test_db_setup
    pass

@pytest.mark.xdist_group(name="filesystem")
def test_file_operations():
    # Different group, may run on different worker
    pass

# In conftest.py
from xdist.scheduler.protocol import Scheduling
from xdist.workermanage import WorkerController

class CustomScheduler:
    """Example custom scheduler implementation."""
    
    def __init__(self, config, log):
        self.config = config
        self.log = log
        self._nodes = []
        self._pending = []
        self._collection_complete = False
    
    @property
    def nodes(self) -> list[WorkerController]:
        return self._nodes
    
    @property 
    def collection_is_completed(self) -> bool:
        return self._collection_complete
    
    def schedule(self) -> None:
        # Custom scheduling logic
        if self.has_pending and self.nodes:
            # Implement custom distribution algorithm
            pass

def pytest_xdist_make_scheduler(config, log):
    if config.getoption("--custom-scheduler"):
        return CustomScheduler(config, log)
    return None  # Use default scheduler

# How pytest-xdist selects schedulers based on --dist option
SCHEDULER_MAP = {
    'each': EachScheduling,
    'load': LoadScheduling, 
    'loadscope': LoadScopeScheduling,
    'loadfile': LoadFileScheduling,
    'loadgroup': LoadGroupScheduling,
    'worksteal': WorkStealingScheduling,
}

def create_scheduler(config, log):
    dist_mode = config.getoption("dist")
    
    # First try custom scheduler hook
    custom = config.hook.pytest_xdist_make_scheduler(config=config, log=log)
    if custom:
        return custom
    
    # Use built-in scheduler
    scheduler_class = SCHEDULER_MAP.get(dist_mode, LoadScheduling)
    return scheduler_class(config, log)

# Choosing the right scheduler for your test suite

# For test suites with:
# - Varying execution times -> LoadScheduling (default)
# - Expensive class/module fixtures -> LoadScopeScheduling  
# - File-based fixtures/setup -> LoadFileScheduling
# - Tests that share state -> LoadGroupScheduling with markers
# - Large test suites -> WorkStealingScheduling
# - Cross-environment testing -> EachScheduling

def pytest_configure(config):
    """Auto-select scheduler based on test suite characteristics."""
    if not config.getoption("dist") or config.getoption("dist") == "no":
        return
        
    # Analyze test suite and suggest optimal scheduler
    test_count = len(config.getoption("file_or_dir") or [])
    
    if test_count > 1000:
        # Large test suite - work stealing might be better
        config.option.dist = "worksteal"
    elif has_expensive_fixtures():
        # Expensive fixtures - use scope grouping
        config.option.dist = "loadscope"