autool-dispatcher/tests/test_remote_worker_controller.py
2026-06-17 19:50:39 +08:00

135 lines
4.7 KiB
Python

import os
import sys
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, PROJECT_ROOT)
from remote_worker_controller import RemoteWorkerController
from config import (
CLEAN_BACKUP_PATH,
LOCAL_IMPORT_DIR,
MUMU_MANAGER_PATH,
MUMU_RECOVER_AHK_EXE,
MUMU_RECOVER_SCRIPT_PATH,
NET_BRIDGE_CARD,
)
def create_controller() -> RemoteWorkerController:
return RemoteWorkerController([])
def test_cmd_task_command_wraps_inner_command_with_cmd_safe_quotes():
command = RemoteWorkerController._cmd_task_command(
'"C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 1 restart'
)
assert (
command
== 'C:\\Windows\\System32\\cmd.exe /d /q /s /c ""C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 1 restart"'
)
def test_cmd_switch_with_command_keeps_original_inner_quotes():
command = RemoteWorkerController._cmd_switch_with_command(
"/c",
'"C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart',
)
assert command == '/c "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart'
def test_restart_mumu_task_command_keeps_schtasks_tr_argument_balanced():
controller = create_controller()
command = controller._restart_mumu_task_command({})
assert isinstance(command, dict)
assert command["command"] == "powershell -NoProfile -ExecutionPolicy Bypass -Command -"
assert "$taskName = 'RestartMuMu';" in command["stdin"]
assert "$taskExecute = 'C:\\Windows\\System32\\cmd.exe';" in command["stdin"]
assert '$taskArguments = \'/c "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart\';' in command["stdin"]
def test_direct_wait_command_quotes_inner_cmd_payload_with_redirection():
controller = create_controller()
command = controller._recover_mumu_command(
{"ssh_target": "192.168.2.71:22", "repo_dir": "D:\\autool"}
)
assert isinstance(command, str)
assert command.startswith("cmd.exe /d /q /s /c ")
assert " & call " in command
assert MUMU_RECOVER_SCRIPT_PATH in command
assert "Recover script not found" in command
assert CLEAN_BACKUP_PATH in command
assert LOCAL_IMPORT_DIR in command
assert MUMU_MANAGER_PATH in command
assert NET_BRIDGE_CARD in command
assert MUMU_RECOVER_AHK_EXE in command
assert "schtasks /create /tn" not in command
def test_start_worker_task_uses_single_wrapped_cmd_argument():
controller = create_controller()
command = controller._start_worker_command(
{
"repo_dir": "C:\\Program Files\\AutoTool Dispatcher",
"python_exe": "C:\\Python39\\python.exe",
}
)
assert isinstance(command, dict)
assert command["command"] == "powershell -NoProfile -ExecutionPolicy Bypass -Command -"
assert "$taskName = 'RunBatch';" in command["stdin"]
assert "$taskExecute = 'C:\\Windows\\System32\\cmd.exe';" in command["stdin"]
assert '$taskArguments = \'/k C:\\Program Files\\AutoTool Dispatcher\\venv\\Scripts\\python.exe C:\\Program Files\\AutoTool Dispatcher\\batch_run.py\';' in command["stdin"]
assert "& schtasks /end /tn $taskName" in command["stdin"]
assert "taskkill /F /T /IM" not in command["stdin"]
def test_stop_worker_uses_simple_taskkill_for_console_hosts():
controller = create_controller()
command = controller._stop_worker_command({})
assert isinstance(command, str)
assert command.startswith("cmd.exe /d /v:on /q /s /c ")
assert 'schtasks /end /tn ""RunBatch""' in command
assert "taskkill /F /T /IM OpenConsole.exe" in command
assert "taskkill /F /T /IM WindowsTerminal.exe" in command
assert "taskkill /F /T /IM conhost.exe" in command
assert "taskkill /F /T /IM powershell.exe" in command
assert "start /min taskkill /F /T /IM cmd.exe" in command
assert "Stop-Process" not in command
assert "Get-CimInstance Win32_Process" not in command
def test_stop_worker_uses_legacy_end_for_protected_workers():
controller = create_controller()
command = controller._stop_worker_command({"worker_id": "192.168.1.51"})
assert isinstance(command, str)
assert 'schtasks /end /tn ""RunBatch""' in command
assert "taskkill /F /T /IM" not in command
def test_start_worker_uses_legacy_end_before_run_for_protected_workers():
controller = create_controller()
command = controller._start_worker_command(
{
"worker_id": "192.168.2.101",
"repo_dir": "C:\\autool",
"python_exe": "C:\\Python39\\python.exe",
}
)
assert isinstance(command, dict)
assert "& schtasks /end /tn $taskName" in command["stdin"]
assert "taskkill /F /T /IM" not in command["stdin"]