OBLITERATUS/tests/test_remote_contracts.py

"""Mutation-safe contracts for remote input and shell construction."""

from __future__ import annotations

import shlex

import pytest

from obliteratus.remote_contracts import (
    normalize_gpu_selection,
    parse_remote_target,
    remote_python_command,
    remote_scp_spec,
    validate_remote_settings,
)


VALID_REMOTE_SETTINGS = {
    "host": "compute.example",
    "user": "runner",
    "port": 22,
    "remote_dir": "/tmp/obliteratus",
    "python": "python3",
    "gpus": None,
    "install_source": "obliteratus",
}


@pytest.mark.parametrize(
    ("raw", "expected"),
    [(None, None), ("all", "all"), ("ALL", "all"), ("0", "0"), (" 02, 0,11 ", "2,0,11")],
)
def test_gpu_selection_normalizes_public_values(raw, expected):
    assert normalize_gpu_selection(raw) == expected


@pytest.mark.parametrize(
    ("raw", "message"),
    [
        ("", "remote gpus must be a non-empty string"),
        ("\x7f", "remote gpus may not contain control characters"),
        ("0,,1", "remote gpus must be 'all' or comma-separated non-negative integers"),
        ("-1", "remote gpus must be 'all' or comma-separated non-negative integers"),
        (
            "0; touch /tmp/pwned",
            "remote gpus must be 'all' or comma-separated non-negative integers",
        ),
        ("gpu0", "remote gpus must be 'all' or comma-separated non-negative integers"),
        ("0\n1", "remote gpus may not contain control characters"),
    ],
)
def test_gpu_selection_rejects_non_device_input(raw, message):
    with pytest.raises(ValueError) as error:
        normalize_gpu_selection(raw)
    assert str(error.value) == message


def test_remote_target_defaults_user_and_parses_single_at_sign():
    assert parse_remote_target("compute.example") == ("root", "compute.example")
    assert parse_remote_target("runner@compute.example") == ("runner", "compute.example")
    assert parse_remote_target(" runner@compute.example ") == ("runner", "compute.example")


@pytest.mark.parametrize(
    ("target", "message"),
    [
        ("", "remote target must be a non-empty string"),
        ("\x7f", "remote target may not contain control characters"),
        ("runner@", "remote host must be a non-empty string"),
        ("@host", "remote user must be a non-empty string"),
        ("bad user@host", "remote user contains unsupported characters"),
        (
            "runner@-oProxyCommand=evil",
            "remote host must be a host name or address without user or options",
        ),
        ("runner@host name", "remote host must be a host name or address without user or options"),
        ("a@b@host", "remote target may contain at most one user separator"),
    ],
)
def test_remote_target_rejects_empty_or_option_like_components(target, message):
    with pytest.raises(ValueError) as error:
        parse_remote_target(target)
    assert str(error.value) == message


@pytest.mark.parametrize(
    ("override", "message"),
    [
        ({"host": ""}, "remote host must be a non-empty string"),
        (
            {"host": "bad@host"},
            "remote host must be a host name or address without user or options",
        ),
        (
            {"host": "bad host"},
            "remote host must be a host name or address without user or options",
        ),
        ({"user": ""}, "remote user must be a non-empty string"),
        ({"user": "-oProxy"}, "remote user contains unsupported characters"),
        ({"user": "bad user"}, "remote user contains unsupported characters"),
        ({"port": 0}, "remote port must be an integer from 1 through 65535"),
        ({"port": 65536}, "remote port must be an integer from 1 through 65535"),
        ({"port": True}, "remote port must be an integer from 1 through 65535"),
        ({"port": "22"}, "remote port must be an integer from 1 through 65535"),
        ({"remote_dir": ""}, "remote directory must be a non-empty string"),
        ({"remote_dir": "relative"}, "remote directory must be an absolute POSIX path"),
        ({"python": ""}, "remote Python must be a non-empty string"),
        ({"install_source": ""}, "remote install source must be a non-empty string"),
        (
            {"install_source": "bad\nsource"},
            "remote install source may not contain control characters",
        ),
    ],
)
def test_remote_settings_reject_invalid_boundaries(override, message):
    values = dict(VALID_REMOTE_SETTINGS)
    values.update(override)
    with pytest.raises(ValueError) as error:
        validate_remote_settings(**values)
    assert str(error.value) == message


@pytest.mark.parametrize("port", [1, 65535])
def test_remote_settings_accept_port_boundaries(port):
    values = dict(VALID_REMOTE_SETTINGS, port=port, gpus=" 02,0 ")
    assert validate_remote_settings(**values) == "2,0"


def test_remote_python_command_preserves_each_untrusted_value_as_one_argument():
    command = remote_python_command(
        "/opt/python builds/current/python",
        ["-m", "obliteratus", "run", "/tmp/a config.yml", "--preset", "x; echo injected"],
        gpus="02,0",
    )
    assert shlex.split(command) == [
        "env",
        "CUDA_VISIBLE_DEVICES=2,0",
        "/opt/python builds/current/python",
        "-m",
        "obliteratus",
        "run",
        "/tmp/a config.yml",
        "--preset",
        "x; echo injected",
    ]


def test_remote_python_command_omits_environment_for_all_devices():
    assert shlex.split(remote_python_command("python3", ["-V"], gpus="all")) == ["python3", "-V"]
    assert shlex.split(remote_python_command("python3", ["-V"])) == ["python3", "-V"]


def test_remote_python_command_reports_invalid_python_contract():
    with pytest.raises(ValueError) as error:
        remote_python_command("", ["-V"])
    assert str(error.value) == "remote Python must be a non-empty string"


def test_remote_scp_spec_quotes_remote_paths_and_directory_suffix():
    assert remote_scp_spec("runner@host", "/tmp/result file", directory=True) == (
        "runner@host:'/tmp/result file/'"
    )
    assert remote_scp_spec("runner@host", "/tmp/results/", directory=True) == (
        "runner@host:/tmp/results/"
    )
    assert remote_scp_spec("runner@host", "/tmp/config.yml") == "runner@host:/tmp/config.yml"


@pytest.mark.parametrize(
    ("target", "path", "message"),
    [
        ("", "/tmp/config.yml", "SSH target must be a non-empty string"),
        ("runner@host", "", "remote SCP path must be a non-empty string"),
        ("runner@host", "/tmp/bad\x7fpath", "remote SCP path may not contain control characters"),
    ],
)
def test_remote_scp_spec_reports_invalid_input_contract(target, path, message):
    with pytest.raises(ValueError) as error:
        remote_scp_spec(target, path)
    assert str(error.value) == message