"""ssh 引数とサーバ上で実行するコマンドの組み立てのテスト(SPEC_03「接続」「取得」「run」)."""
from __future__ import annotations
import pytest
from tias_judge_model.remote.commands import (
build_job_env,
build_probe_command,
build_ssh_argv,
validate_job_name,
)
def _argv(jump: str | None, timeout: int = 15) -> list[str]:
return build_ssh_argv(
"ssh",
user="hasegawa",
address="192.168.23.14",
jump_host=jump,
connect_timeout_sec=timeout,
remote_command="echo hi",
)
def _has_option(argv: list[str], value: str) -> bool:
return any(argv[i] == "-o" and argv[i + 1] == value for i in range(len(argv) - 1))
class TestBuildSshArgv:
def test_with_jump_host(self) -> None:
argv = _argv("labpc")
assert argv[0] == "ssh"
j = argv.index("-J")
assert argv[j + 1] == "labpc"
assert "hasegawa@192.168.23.14" in argv
assert argv[-1] == "echo hi"
# 接続先より前にオプションがある
assert j < argv.index("hasegawa@192.168.23.14")
def test_without_jump_host(self) -> None:
argv = _argv(None)
assert "-J" not in argv
assert "hasegawa@192.168.23.14" in argv
@pytest.mark.parametrize("jump", ["labpc", None])
def test_batch_mode_and_timeout_always(self, jump: str | None) -> None:
argv = _argv(jump, timeout=7)
assert _has_option(argv, "BatchMode=yes")
assert _has_option(argv, "ConnectTimeout=7")
def test_remote_command_is_single_argument(self) -> None:
argv = build_ssh_argv(
"ssh",
user="u",
address="h",
jump_host=None,
connect_timeout_sec=5,
remote_command="cd ~/x; echo a b",
)
assert argv[-1] == "cd ~/x; echo a b"
@pytest.mark.parametrize("bad", ["-oProxyCommand=evil", "a b", ""])
def test_rejects_option_like_destination(self, bad: str) -> None:
with pytest.raises(ValueError):
build_ssh_argv(
"ssh",
user="u",
address=bad,
jump_host=None,
connect_timeout_sec=5,
remote_command="true",
)
class TestBuildProbeCommand:
def test_queries_listed_in_spec(self) -> None:
command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500)
assert "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu" in command
assert "--query-compute-apps=gpu_uuid,pid,used_memory" in command
assert "--format=csv,noheader,nounits" in command
def test_observation_uses_interval_and_duration(self) -> None:
command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500)
assert "-lms 500" in command
assert "timeout 10" in command
assert "pmon -s u" in command
def test_observe_zero_does_not_observe(self) -> None:
command = build_probe_command("~/work/x", observe_sec=0, sample_interval_ms=500)
assert "-lms" not in command
assert "pmon" not in command
assert "--query-gpu=index,uuid,name" in command
def test_is_single_line(self) -> None:
command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500)
assert "\n" not in command
class TestJobSettings:
@pytest.mark.parametrize("name", ["a", "exp1", "Exp_1-b", "x" * 64])
def test_valid_job_names(self, name: str) -> None:
assert validate_job_name(name) == name
@pytest.mark.parametrize("name", ["", "_a", "-a", "a b", "a.b", "a:b", "x" * 65, "a/b"])
def test_invalid_job_names(self, name: str) -> None:
with pytest.raises(ValueError):
validate_job_name(name)
def test_job_env_contains_selected_gpu_and_extra(self) -> None:
env = build_job_env(2, ["FOO=bar", "EMPTY="])
assert env["CUDA_VISIBLE_DEVICES"] == "2"
assert env["FOO"] == "bar"
assert env["EMPTY"] == ""