86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
|
"""Mock transport that runs workers locally in isolated temp directories.
|
||
|
|
|
||
|
|
Each MockConnection gets its own temp dir standing in for /workspace/.
|
||
|
|
push/pull copy files into/out of that dir. exec runs commands as local
|
||
|
|
subprocesses with /workspace paths rewritten to the temp dir.
|
||
|
|
|
||
|
|
Supports injecting failures: call conn.kill() to simulate a node dying.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
|
||
|
|
class MockConnection:
|
||
|
|
def __init__(self, root_dir: str):
|
||
|
|
self.root_dir = root_dir
|
||
|
|
self._alive = True
|
||
|
|
|
||
|
|
async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]:
|
||
|
|
if not self._alive:
|
||
|
|
raise ConnectionError("node is dead")
|
||
|
|
|
||
|
|
cmd = cmd.replace("/workspace", self.root_dir)
|
||
|
|
|
||
|
|
proc = await asyncio.create_subprocess_shell(
|
||
|
|
cmd,
|
||
|
|
cwd=self.root_dir,
|
||
|
|
stdout=asyncio.subprocess.PIPE,
|
||
|
|
stderr=asyncio.subprocess.PIPE,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
stdout, stderr = await asyncio.wait_for(
|
||
|
|
proc.communicate(), timeout=timeout
|
||
|
|
)
|
||
|
|
return proc.returncode, (stdout + stderr).decode()
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
proc.kill()
|
||
|
|
return 1, "timeout"
|
||
|
|
|
||
|
|
async def push(self, local_path: str, remote_path: str) -> None:
|
||
|
|
if not self._alive:
|
||
|
|
raise ConnectionError("node is dead")
|
||
|
|
dest = remote_path.replace("/workspace", self.root_dir)
|
||
|
|
os.makedirs(os.path.dirname(dest) or self.root_dir, exist_ok=True)
|
||
|
|
shutil.copy2(local_path, dest)
|
||
|
|
|
||
|
|
async def pull(self, remote_path: str, local_path: str) -> None:
|
||
|
|
if not self._alive:
|
||
|
|
raise ConnectionError("node is dead")
|
||
|
|
src = remote_path.replace("/workspace", self.root_dir)
|
||
|
|
if not os.path.exists(src):
|
||
|
|
raise FileNotFoundError(f"remote file not found: {src}")
|
||
|
|
os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
|
||
|
|
shutil.copy2(src, local_path)
|
||
|
|
|
||
|
|
async def close(self) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def kill(self):
|
||
|
|
"""Simulate node death."""
|
||
|
|
self._alive = False
|
||
|
|
|
||
|
|
def revive(self):
|
||
|
|
"""Bring node back (for testing reconnect scenarios)."""
|
||
|
|
self._alive = True
|
||
|
|
|
||
|
|
|
||
|
|
class MockTransport:
|
||
|
|
"""Creates MockConnections backed by temp directories."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._dirs: list[str] = []
|
||
|
|
|
||
|
|
async def connect(self, host: str, port: int) -> MockConnection:
|
||
|
|
d = tempfile.mkdtemp(prefix=f"mock_node_{host}_{port}_")
|
||
|
|
self._dirs.append(d)
|
||
|
|
return MockConnection(d)
|
||
|
|
|
||
|
|
def cleanup(self):
|
||
|
|
for d in self._dirs:
|
||
|
|
shutil.rmtree(d, ignore_errors=True)
|