grm/scripts/ci/distribute_molecule.py

#!/usr/bin/env python3
"""Distribute molecule (scenario, platform) pairs across N parallel runners.

Discovers all molecule scenarios under ansible/roles/*/molecule/ and
crosses them with the supported OS platform matrix, then splits the
resulting test pairs evenly across the requested number of runners.

Each pair is printed as ``scenario|platform_name|platform_image|platform_command``
so the CI workflow can set the appropriate environment variables.

Usage:
  python3 scripts/distribute_molecule.py --runner-index 1 --max-runners 3
  # prints: default|ubuntu-2204|ubuntu:22.04| lifecycle|ubuntu-2204|ubuntu:22.04| ...
  python3 scripts/distribute_molecule.py --list
  # prints all scenarios, one per line
  python3 scripts/distribute_molecule.py --list-platforms
  # prints all platforms, one per line
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import click

from gitea_runner_manager.i18n import _

DEFAULT_MAX_RUNNERS = 3
MOLECULE_ROOT = Path("ansible/roles/gitea-runner/molecule")

#: Supported OS platform matrix.
#: Each entry maps a short name to (image, command).
#: The command must be systemd since rootless Docker requires loginctl/systemctl --user.
PLATFORMS: list[dict[str, str]] = [
    {"name": "ubuntu-2204", "image": "geerlingguy/docker-ubuntu2204-ansible:latest", "command": "/lib/systemd/systemd"},
    {"name": "ubuntu-2404", "image": "geerlingguy/docker-ubuntu2404-ansible:latest", "command": "/lib/systemd/systemd"},
    {"name": "debian-12", "image": "geerlingguy/docker-debian12-ansible:latest", "command": "/lib/systemd/systemd"},
    {"name": "archlinux", "image": "marcstraube/archlinux-ansible:latest", "command": "/usr/lib/systemd/systemd"},
]


@dataclass(frozen=True)
class TestPair:
    """A (scenario, platform) combination to test."""

    scenario: str
    platform: dict[str, str]

    def encode(self) -> str:
        """Serialize to a pipe-delimited string for CI consumption."""
        return f"{self.scenario}|{self.platform['name']}|{self.platform['image']}|{self.platform['command']}"

    @staticmethod
    def decode(encoded: str) -> TestPair:
        """Deserialize from a pipe-delimited string."""
        parts = encoded.split("|")
        return TestPair(
            scenario=parts[0],
            platform={"name": parts[1], "image": parts[2], "command": parts[3]},
        )


def discover_scenarios(root: Path | None = None) -> list[str]:
    """Return sorted list of molecule scenario directory names."""
    if root is None:
        root = MOLECULE_ROOT
    if not root.is_dir():
        raise click.ClickException(_("Molecule directory not found: {path}", path=str(root)))
    scenarios = [d.name for d in root.iterdir() if d.is_dir() and not d.name.startswith("_") and d.name != "common"]
    return sorted(scenarios)


def build_pairs(scenarios: list[str], platforms: list[dict[str, str]] | None = None) -> list[TestPair]:
    """Build the full cross-product of scenarios and platforms."""
    if platforms is None:
        platforms = PLATFORMS
    return [TestPair(s, p) for s in scenarios for p in platforms]


def distribute(pairs: list[TestPair], max_runners: int) -> list[list[TestPair]]:
    """Split *pairs* into *max_runners* balanced groups (round-robin)."""
    groups: list[list[TestPair]] = [[] for _ in range(max_runners)]
    for i, pair in enumerate(pairs):
        groups[i % max_runners].append(pair)
    return groups


def pairs_for_runner(pairs: list[TestPair], runner_index: int, max_runners: int) -> list[TestPair]:
    """Return the subset of pairs assigned to *runner_index*."""
    groups = distribute(pairs, max_runners)
    if runner_index < 0 or runner_index >= len(groups):
        raise click.ClickException(
            _(
                "Runner index {index} out of range (0..{max})",
                index=runner_index,
                max=max_runners - 1,
            )
        )
    return groups[runner_index]


def _write_github_env(key: str, value: str) -> None:
    """Append a key=value line to the $GITHUB_ENV file."""
    import os

    gh_env = os.environ.get("GITHUB_ENV")
    if not gh_env:
        raise click.ClickException("GITHUB_ENV environment variable is not set")
    with open(gh_env, "a") as f:  # noqa: PTH123
        f.write(f"{key}={value}\n")


@click.command()
@click.option(
    "--runner-index",
    type=int,
    default=None,
    help="One-based runner index (Gitea Actions renders 0 as empty). "
    "Converted to zero-based internally. If omitted, prints all groups.",
)
@click.option(
    "--max-runners",
    type=int,
    default=DEFAULT_MAX_RUNNERS,
    show_default=True,
    help="Total number of parallel runners.",
)
@click.option(
    "--list",
    "list_all",
    is_flag=True,
    help="List all discovered scenarios, one per line.",
)
@click.option(
    "--list-platforms",
    "list_platforms",
    is_flag=True,
    help="List all supported platforms, one per line.",
)
@click.option(
    "--github-env",
    "github_env",
    is_flag=True,
    default=False,
    help="Write TEST_PAIRS and SKIP to $GITHUB_ENV (for CI workflow steps).",
)
@click.option(
    "--skip-if-excess",
    is_flag=True,
    default=False,
    help="With --github-env: write SKIP=true when runner-index exceeds max-runners.",
)
def cli(
    runner_index: int | None,
    max_runners: int,
    list_all: bool,
    list_platforms: bool,
    github_env: bool,
    skip_if_excess: bool,
) -> None:
    scenarios = discover_scenarios()
    if list_all:
        for s in scenarios:
            click.echo(s)
        return
    if list_platforms:
        for p in PLATFORMS:
            click.echo(f"{p['name']}|{p['image']}|{p['command']}")
        return
    pairs = build_pairs(scenarios)
    if runner_index is None:
        groups = distribute(pairs, max_runners)
        for i, group in enumerate(groups):
            labels = " ".join(p.encode() for p in group) if group else "(none)"
            click.echo(f"Runner {i}: {labels}")
        return

    # Skip if runner index exceeds available runners (CI static matrix has 3 slots)
    if skip_if_excess and github_env and runner_index > max_runners:
        click.echo(f"Skipping — runner index {runner_index} > max runners {max_runners}")
        _write_github_env("TEST_PAIRS", "")
        _write_github_env("SKIP", "true")
        return

    # Convert 1-based CLI index to 0-based internal index
    zero_based = runner_index - 1
    assigned = pairs_for_runner(pairs, zero_based, max_runners)
    encoded = " ".join(p.encode() for p in assigned)

    if github_env:
        _write_github_env("TEST_PAIRS", encoded)
        _write_github_env("SKIP", "false")
        click.echo(f"Assigned pairs: {encoded}")
        return

    click.echo(encoded)


if __name__ == "__main__":  # pragma: no cover
    cli()  # pragma: no cover