grm/scripts/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 0 --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]


@click.command()
@click.option(
    "--runner-index",
    type=int,
    default=None,
    help="Zero-based runner index. 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.",
)
def cli(runner_index: int | None, max_runners: int, list_all: bool, list_platforms: 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
    assigned = pairs_for_runner(pairs, runner_index, max_runners)
    click.echo(" ".join(p.encode() for p in assigned))


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