pytest-balance 0.1.0__py3-none-any.whl

pytest_balance/__init__.py

@@ -0,0 +1,3 @@
+"""Duration-based test distribution for pytest."""
+
+__version__ = "0.1.0"

pytest_balance/__main__.py

@@ -0,0 +1,5 @@
+"""Allow running pytest-balance as a module: python -m pytest_balance"""
+
+from pytest_balance.cli import main
+
+main()

pytest_balance/algorithms/lpt.py

@@ -0,0 +1,68 @@
+"""Longest Processing Time First (LPT) partitioning algorithm."""
+
+from __future__ import annotations
+
+import heapq
+
+from pytest_balance.algorithms.partitioner import Scope, group_by_scope
+from pytest_balance.store.models import DurationEstimate
+from pytest_balance.store.reader import default_estimate
+
+
+def partition(durations: dict[str, float], n: int) -> list[list[str]]:
+    """Partition items into n buckets minimizing makespan using LPT.
+
+    Args:
+        durations: Mapping of item ID to estimated duration.
+        n: Number of buckets.
+
+    Returns:
+        List of n lists, each containing item IDs assigned to that bucket.
+    """
+    if n <= 0:
+        raise ValueError(f"n must be >= 1, got {n}")
+
+    buckets: list[list[str]] = [[] for _ in range(n)]
+
+    if not durations:
+        return buckets
+
+    # Sort by duration descending, break ties by name for determinism
+    sorted_items = sorted(durations.items(), key=lambda x: (-x[1], x[0]))
+
+    # Min-heap of (total_duration, bucket_index)
+    heap: list[tuple[float, int]] = [(0.0, i) for i in range(n)]
+    heapq.heapify(heap)
+
+    for item_id, _ in sorted_items:
+        total, idx = heapq.heappop(heap)
+        buckets[idx].append(item_id)
+        heapq.heappush(heap, (total + durations[item_id], idx))
+
+    return buckets
+
+
+def compute_order(
+    collection: list[str],
+    estimates: dict[str, DurationEstimate],
+    scope: Scope,
+) -> list[int]:
+    """Return indices of `collection` ordered LPT scope-adjacent.
+
+    Groups are built by scope, sorted by descending estimated duration with a
+    lexicographic tie-break on scope_id, and tests of the same group are
+    emitted consecutively. Pure function: deterministic, no side effects.
+    """
+    if not collection:
+        return []
+
+    fallback = default_estimate(estimates)
+    groups = group_by_scope(collection, scope)
+    for group in groups:
+        group.estimated_duration = sum(
+            estimates.get(tid, fallback).estimate for tid in group.test_ids
+        )
+    groups.sort(key=lambda g: (-g.estimated_duration, g.scope_id))
+
+    test_id_to_index = {tid: idx for idx, tid in enumerate(collection)}
+    return [test_id_to_index[tid] for group in groups for tid in group.test_ids]

pytest_balance/algorithms/partitioner.py

@@ -0,0 +1,62 @@
+"""Scope-aware test grouping for balanced partitioning."""
+
+from __future__ import annotations
+
+import enum
+from collections import OrderedDict
+from dataclasses import dataclass, field
+
+
+class Scope(enum.Enum):
+    TEST = "test"
+    CLASS = "class"
+    MODULE = "module"
+    GROUP = "group"
+
+
+@dataclass
+class TestGroup:
+    scope_id: str
+    test_ids: list[str] = field(default_factory=list)
+    estimated_duration: float = 0.0
+
+
+def extract_scope(test_id: str, scope: Scope) -> str:
+    """Extract the scope key from a pytest node ID."""
+    if scope == Scope.TEST:
+        return test_id
+
+    if scope == Scope.GROUP:
+        bracket_pos = test_id.rfind("]")
+        at_pos = test_id.rfind("@")
+        if at_pos > bracket_pos:
+            return test_id[at_pos + 1 :]
+        return test_id
+
+    parts = test_id.split("::")
+    file_part = parts[0] if parts else test_id
+
+    if scope == Scope.MODULE:
+        return file_part
+
+    if scope == Scope.CLASS:
+        if len(parts) >= 3:
+            return f"{parts[0]}::{parts[1]}"
+        return test_id
+
+    return test_id
+
+
+def group_by_scope(test_ids: list[str], scope: Scope) -> list[TestGroup]:
+    """Group test IDs by scope, preserving order within groups."""
+    if not test_ids:
+        return []
+
+    groups: OrderedDict[str, TestGroup] = OrderedDict()
+    for test_id in test_ids:
+        scope_key = extract_scope(test_id, scope)
+        if scope_key not in groups:
+            groups[scope_key] = TestGroup(scope_id=scope_key)
+        groups[scope_key].test_ids.append(test_id)
+
+    return list(groups.values())

pytest_balance/ci/detect.py

@@ -0,0 +1,209 @@
+"""CI environment auto-detection."""
+
+from __future__ import annotations
+
+import os
+import warnings
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class CIContext:
+    provider: str
+    node_index: int
+    node_total: int
+    run_id: str
+    branch: str | None = None
+
+
+def detect_ci(
+    explicit_index: int | None = None,
+    explicit_total: int | None = None,
+) -> CIContext | None:
+    ctx = _detect_provider()
+
+    if explicit_index is not None and explicit_total is not None:
+        if ctx is not None:
+            ctx = CIContext(
+                provider=ctx.provider,
+                node_index=explicit_index,
+                node_total=explicit_total,
+                run_id=ctx.run_id,
+                branch=ctx.branch,
+            )
+        else:
+            ctx = CIContext(
+                provider="explicit",
+                node_index=explicit_index,
+                node_total=explicit_total,
+                run_id="local",
+            )
+
+    if ctx is None:
+        return None
+
+    _validate(ctx)
+    return ctx
+
+
+def _detect_provider() -> CIContext | None:
+    for detector in [
+        _detect_github,
+        _detect_gitlab,
+        _detect_circleci,
+        _detect_azure,
+        _detect_buildkite,
+        _detect_generic,
+    ]:
+        ctx = detector()
+        if ctx is not None:
+            return ctx
+    return None
+
+
+def _detect_github() -> CIContext | None:
+    if os.environ.get("GITHUB_ACTIONS") != "true":
+        return None
+    index = os.environ.get("PYTEST_BALANCE_NODE_INDEX")
+    total = os.environ.get("PYTEST_BALANCE_NODE_TOTAL")
+    if index is None or total is None:
+        return None
+    try:
+        index_int, total_int = int(index), int(total)
+    except ValueError:
+        warnings.warn(
+            f"Non-integer balance env var: NODE_INDEX={index!r}, NODE_TOTAL={total!r}",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+    run_id = os.environ.get("GITHUB_RUN_ID", "unknown")
+    attempt = os.environ.get("GITHUB_RUN_ATTEMPT", "1")
+    branch = os.environ.get("GITHUB_REF_NAME")
+    return CIContext("github", index_int, total_int, f"{run_id}-{attempt}", branch)
+
+
+def _detect_gitlab() -> CIContext | None:
+    if os.environ.get("GITLAB_CI") != "true":
+        return None
+    index = os.environ.get("CI_NODE_INDEX")
+    total = os.environ.get("CI_NODE_TOTAL")
+    if index is None or total is None:
+        return None
+    try:
+        index_int, total_int = int(index) - 1, int(total)
+    except ValueError:
+        warnings.warn(
+            f"Non-integer CI env var: CI_NODE_INDEX={index!r}, CI_NODE_TOTAL={total!r}",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+    return CIContext(
+        "gitlab",
+        index_int,
+        total_int,
+        os.environ.get("CI_PIPELINE_ID", "unknown"),
+        os.environ.get("CI_COMMIT_REF_NAME"),
+    )
+
+
+def _detect_circleci() -> CIContext | None:
+    if os.environ.get("CIRCLECI") != "true":
+        return None
+    index = os.environ.get("CIRCLE_NODE_INDEX")
+    total = os.environ.get("CIRCLE_NODE_TOTAL")
+    if index is None or total is None:
+        return None
+    try:
+        index_int, total_int = int(index), int(total)
+    except ValueError:
+        warnings.warn(
+            f"Non-integer CI env var: CIRCLE_NODE_INDEX={index!r}, CIRCLE_NODE_TOTAL={total!r}",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+    return CIContext(
+        "circleci",
+        index_int,
+        total_int,
+        os.environ.get("CIRCLE_BUILD_NUM", "unknown"),
+        os.environ.get("CIRCLE_BRANCH"),
+    )
+
+
+def _detect_azure() -> CIContext | None:
+    if os.environ.get("TF_BUILD") != "True":
+        return None
+    index = os.environ.get("SYSTEM_JOBPOSITIONINPHASE")
+    total = os.environ.get("SYSTEM_TOTALJOBSINPHASE")
+    if index is None or total is None:
+        return None
+    try:
+        index_int, total_int = int(index) - 1, int(total)
+    except ValueError:
+        warnings.warn(
+            f"Non-integer CI env var: SYSTEM_JOBPOSITIONINPHASE={index!r},"
+            f" SYSTEM_TOTALJOBSINPHASE={total!r}",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+    return CIContext(
+        "azure",
+        index_int,
+        total_int,
+        os.environ.get("BUILD_BUILDID", "unknown"),
+        os.environ.get("BUILD_SOURCEBRANCH"),
+    )
+
+
+def _detect_buildkite() -> CIContext | None:
+    if os.environ.get("BUILDKITE") != "true":
+        return None
+    index = os.environ.get("BUILDKITE_PARALLEL_JOB")
+    total = os.environ.get("BUILDKITE_PARALLEL_JOB_COUNT")
+    if index is None or total is None:
+        return None
+    try:
+        index_int, total_int = int(index), int(total)
+    except ValueError:
+        warnings.warn(
+            f"Non-integer CI env var: BUILDKITE_PARALLEL_JOB={index!r},"
+            f" BUILDKITE_PARALLEL_JOB_COUNT={total!r}",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+    return CIContext(
+        "buildkite",
+        index_int,
+        total_int,
+        os.environ.get("BUILDKITE_BUILD_ID", "unknown"),
+        os.environ.get("BUILDKITE_BRANCH"),
+    )
+
+
+def _detect_generic() -> CIContext | None:
+    index = os.environ.get("PYTEST_BALANCE_NODE_INDEX")
+    total = os.environ.get("PYTEST_BALANCE_NODE_TOTAL")
+    if index is None or total is None:
+        return None
+    try:
+        index_int, total_int = int(index), int(total)
+    except ValueError:
+        warnings.warn(
+            f"Non-integer balance env var: NODE_INDEX={index!r}, NODE_TOTAL={total!r}",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+    return CIContext("unknown", index_int, total_int, "generic")
+
+
+def _validate(ctx: CIContext) -> None:
+    if ctx.node_total < 1:
+        raise ValueError(f"node_total must be >= 1, got {ctx.node_total}")
+    if not (0 <= ctx.node_index < ctx.node_total):
+        raise ValueError(f"node_index {ctx.node_index} out of range for {ctx.node_total} nodes")

pytest_balance/ci/splitter.py

@@ -0,0 +1,46 @@
+"""CI-level test splitting using LPT and scope-aware grouping."""
+
+from __future__ import annotations
+
+from pytest_balance.algorithms.lpt import partition
+from pytest_balance.algorithms.partitioner import Scope, group_by_scope
+from pytest_balance.store.models import DurationEstimate
+from pytest_balance.store.reader import default_estimate
+
+
+def split_tests(
+    tests: list[str],
+    estimates: dict[str, DurationEstimate],
+    node_index: int,
+    node_total: int,
+    scope: Scope,
+) -> list[str]:
+    if not tests:
+        return []
+
+    fallback = default_estimate(estimates)
+
+    # Group tests by scope
+    groups = group_by_scope(tests, scope)
+
+    # Compute estimated duration for each group
+    group_durations: dict[str, float] = {}
+    for group in groups:
+        group.estimated_duration = sum(
+            estimates.get(tid, fallback).estimate for tid in group.test_ids
+        )
+        group_durations[group.scope_id] = group.estimated_duration
+
+    # LPT partition groups across nodes
+    buckets = partition(group_durations, node_total)
+
+    # Select groups for this node
+    selected_scopes = set(buckets[node_index])
+
+    # Flatten test IDs preserving original order within groups
+    selected_tests: list[str] = []
+    for group in groups:
+        if group.scope_id in selected_scopes:
+            selected_tests.extend(group.test_ids)
+
+    return selected_tests