devx/tests/unit/test_distribute_files.py

"""Unit tests for devx.ci.distribute_files."""

from pathlib import Path

import pytest
from click.testing import CliRunner

from devx.ci.distribute_files import (
    DEFAULT_MAX_RUNNERS,
    _file_weight,
    discover_files,
    distribute,
    files_for_runner,
    main,
)


class TestDiscoverFiles:
    def test_discovers_sorted(self, tmp_path: Path) -> None:
        (tmp_path / "test_b.py").write_text("")
        (tmp_path / "test_a.py").write_text("")
        result = discover_files(str(tmp_path / "test_*.py"))
        assert len(result) == 2
        assert result[0].endswith("test_a.py")
        assert result[1].endswith("test_b.py")

    def test_no_matches(self, tmp_path: Path) -> None:
        assert discover_files(str(tmp_path / "nonexistent-*.py")) == []


class TestDistribute:
    def test_even_split(self) -> None:
        files = [f"test_{i}.py" for i in range(6)]
        groups = distribute(files, 3)
        assert len(groups) == 3
        assert all(len(g) == 2 for g in groups)

    def test_uneven_split(self) -> None:
        files = [f"test_{i}.py" for i in range(5)]
        groups = distribute(files, 3)
        assert len(groups[0]) == 2
        assert len(groups[1]) == 2
        assert len(groups[2]) == 1

    def test_more_runners_than_files(self) -> None:
        files = ["test_a.py"]
        groups = distribute(files, 5)
        assert len(groups) == 5
        assert len(groups[0]) == 1
        assert all(len(g) == 0 for g in groups[1:])

    def test_empty(self) -> None:
        assert distribute([], 3) == [[], [], []]


class TestFilesForRunner:
    def test_returns_correct_subset(self) -> None:
        files = [f"test_{i}.py" for i in range(6)]
        assert len(files_for_runner(files, 0, 3)) == 2
        assert len(files_for_runner(files, 1, 3)) == 2
        assert len(files_for_runner(files, 2, 3)) == 2

    def test_out_of_range_raises(self) -> None:
        with pytest.raises(Exception, match="out of range"):
            files_for_runner(["a.py"], 5, 3)


class TestCli:
    def test_no_runner_index_prints_groups(self, tmp_path: Path) -> None:
        for i in range(3):
            (tmp_path / f"test_{i}.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(main, ["--pattern", str(tmp_path / "test_*.py"), "--max-runners", "3"])
        assert result.exit_code == 0
        assert "Runner 0:" in result.output
        assert "Runner 1:" in result.output
        assert "Runner 2:" in result.output

    def test_runner_index_prints_assigned(self, tmp_path: Path) -> None:
        for i in range(3):
            (tmp_path / f"test_{i}.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["--pattern", str(tmp_path / "test_*.py"), "--runner-index", "1", "--max-runners", "3"],
        )
        assert result.exit_code == 0
        assert "test_0.py" in result.output

    def test_github_env_writes_files(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        gh_file = tmp_path / "env.txt"
        monkeypatch.setenv("GITHUB_ENV", str(gh_file))
        for i in range(2):
            (tmp_path / f"test_{i}.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["--pattern", str(tmp_path / "test_*.py"), "--runner-index", "1", "--max-runners", "2", "--github-env"],
        )
        assert result.exit_code == 0
        content = gh_file.read_text()
        assert "ASSIGNED_FILES=" in content
        assert "SKIP=false" in content

    def test_github_env_multiline_uses_heredoc(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        gh_file = tmp_path / "env.txt"
        monkeypatch.setenv("GITHUB_ENV", str(gh_file))
        for i in range(6):
            (tmp_path / f"test_{i}.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["--pattern", str(tmp_path / "test_*.py"), "--runner-index", "1", "--max-runners", "2", "--github-env"],
        )
        assert result.exit_code == 0
        content = gh_file.read_text()
        # Multi-line values must use heredoc syntax to avoid corrupting $GITHUB_ENV
        assert "ASSIGNED_FILES<<EOF" in content
        assert content.count("EOF") >= 2
        assert "SKIP=false" in content

    def test_skip_if_excess(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        gh_file = tmp_path / "env.txt"
        monkeypatch.setenv("GITHUB_ENV", str(gh_file))
        (tmp_path / "test.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(
            main,
            [
                "--pattern",
                str(tmp_path / "test_*.py"),
                "--runner-index",
                "5",
                "--max-runners",
                "2",
                "--github-env",
                "--skip-if-excess",
            ],
        )
        assert result.exit_code == 0
        content = gh_file.read_text()
        assert "ASSIGNED_FILES=\n" in content
        assert "SKIP=true" in content

    def test_runner_index_zero_raises(self, tmp_path: Path) -> None:
        (tmp_path / "test.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["--pattern", str(tmp_path / "test_*.py"), "--runner-index", "0", "--max-runners", "3"],
        )
        assert result.exit_code != 0

    def test_no_env_var_raises(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.delenv("GITHUB_ENV", raising=False)
        (tmp_path / "test.py").write_text("")
        runner = CliRunner()
        result = runner.invoke(
            main,
            ["--pattern", str(tmp_path / "test_*.py"), "--runner-index", "1", "--max-runners", "3", "--github-env"],
        )
        assert result.exit_code != 0


def test_default_max_runners() -> None:
    assert DEFAULT_MAX_RUNNERS == 3


def test_main_module_block() -> None:
    import devx.ci.distribute_files as mod

    assert hasattr(mod, "main")


class TestFileWeight:
    def test_weight_based_on_size(self, tmp_path: Path) -> None:
        f = tmp_path / "test_big.py"
        f.write_text("x" * 5000)
        assert _file_weight(str(f)) == 5000

    def test_min_weight_is_1(self, tmp_path: Path) -> None:
        f = tmp_path / "empty.py"
        f.write_text("")
        assert _file_weight(str(f)) == 1

    def test_nonexistent_file_returns_1(self) -> None:
        assert _file_weight("/nonexistent/file.py") == 1


class TestDistributeLpt:
    def test_large_files_on_different_runners(self, tmp_path: Path) -> None:
        """Two large files should go to different runners."""
        big1 = tmp_path / "test_big1.py"
        big2 = tmp_path / "test_big2.py"
        small1 = tmp_path / "test_small1.py"
        small2 = tmp_path / "test_small2.py"
        big1.write_text("x" * 10000)
        big2.write_text("x" * 10000)
        small1.write_text("x")
        small2.write_text("x")
        files = [str(big1), str(big2), str(small1), str(small2)]
        groups = distribute(files, 2)
        runner_0 = groups[0]
        runner_1 = groups[1]
        # Big files should be on different runners
        assert not (str(big1) in runner_0 and str(big2) in runner_0)
        assert not (str(big1) in runner_1 and str(big2) in runner_1)

    def test_all_files_preserved(self, tmp_path: Path) -> None:
        for i in range(5):
            (tmp_path / f"test_{i}.py").write_text(f"content {i}" * (i + 1))
        files = [str(tmp_path / f"test_{i}.py") for i in range(5)]
        groups = distribute(files, 3)
        flat = sorted(f for group in groups for f in group)
        assert flat == sorted(files)