166 lines
6.8 KiB
Python
166 lines
6.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import pwd
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from tests.test_guest_exec import FakeScope
|
|
from tests.test_lifecycle import LifecycleTestCase
|
|
from uvm import guest_agent as agent
|
|
from uvm.errors import ConflictError, GuestAgentUnavailableError, NotFoundError
|
|
from uvm.guest_terminal import GuestTerminal
|
|
|
|
|
|
class GuestTerminalTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.scope = FakeScope()
|
|
self.addCleanup(self.scope.temporary.cleanup)
|
|
self.home = self.enterContext(tempfile.TemporaryDirectory())
|
|
username = pwd.getpwuid(os.geteuid()).pw_name
|
|
original_popen = subprocess.Popen
|
|
# Exercise the actual guest child entry point with a temporary account
|
|
# home, so interactive shells never read/write the developer's dotfiles.
|
|
script = f"""
|
|
import os, runpy
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
agent = runpy.run_path({agent.__file__!r})
|
|
agent['pwd'].getpwnam = lambda _: SimpleNamespace(pw_name={username!r}, pw_uid=os.geteuid(),
|
|
pw_gid=os.getegid(), pw_dir={self.home!r}, pw_shell='/bin/sh')
|
|
agent['child_main'](Path({str(self.scope.path)!r}), {username!r}, None)
|
|
"""
|
|
|
|
def launch(_args, **kwargs):
|
|
return original_popen([sys.executable, "-I", "-c", script], **kwargs)
|
|
|
|
self.enterContext(patch.object(agent.subprocess, "Popen", side_effect=launch))
|
|
left, right = socket.socketpair()
|
|
self.terminal = GuestTerminal(left)
|
|
self.addCleanup(self.terminal.close)
|
|
|
|
def runner(connection, account, rows, cols):
|
|
return agent.run_terminal(connection, account, rows, cols, scope_factory=lambda: self.scope)
|
|
|
|
self.thread = threading.Thread(target=agent.serve_connection,
|
|
args=(right, {"vm_id": "vm-test", "username": username}),
|
|
kwargs={"terminal_runner": runner}, daemon=True)
|
|
self.thread.start()
|
|
self.addCleanup(self.thread.join, 3)
|
|
agent.receive_frame(left, 4096, time.monotonic() + 2)
|
|
await asyncio.to_thread(self.terminal.start, 24, 80)
|
|
self.prompt = b"# " if os.geteuid() == 0 else b"$ "
|
|
await self.until(self.prompt)
|
|
|
|
async def until(self, wanted):
|
|
data = bytearray()
|
|
try:
|
|
async with asyncio.timeout(5):
|
|
while wanted not in data:
|
|
message = await self.terminal.read()
|
|
self.assertEqual(message["type"], "output", message)
|
|
data.extend(message["data"])
|
|
except TimeoutError:
|
|
self.fail(f"Missing {wanted!r} in terminal output: {bytes(data)!r}")
|
|
return bytes(data)
|
|
|
|
async def asyncTearDown(self):
|
|
self.terminal.close()
|
|
await asyncio.to_thread(self.thread.join, 3)
|
|
self.assertFalse(self.thread.is_alive())
|
|
|
|
async def test_terminal_has_tty_persistent_shell_state_and_exit_code(self):
|
|
await self.terminal.write(b"test -t 0 && printf '\\nREAL-TTY\\n'\r")
|
|
await self.until(b"\r\nREAL-TTY\r\n")
|
|
await self.terminal.write(b"export NEST_TEST=hello; cd /\r")
|
|
await self.terminal.write(b"printf 'STATE:%s:%s\\n' \"$NEST_TEST\" \"$PWD\"\r")
|
|
await self.until(b"STATE:hello:/\r\n")
|
|
await self.terminal.write(b"exit 7\r")
|
|
async with asyncio.timeout(5):
|
|
while True:
|
|
message = await self.terminal.read()
|
|
if message["type"] == "exit":
|
|
self.assertEqual(message["exit_code"], 7)
|
|
break
|
|
|
|
async def test_resize_and_control_c_reach_the_guest_foreground_job(self):
|
|
await self.terminal.resize(40, 120)
|
|
await self.terminal.write(b"stty size\r")
|
|
await self.until(b"40 120\r\n")
|
|
await self.terminal.write(b"sh -c 'printf \"\\nJOB-STARTED\\n\"; exec sleep 20'\r")
|
|
await self.until(b"\r\nJOB-STARTED\r\n")
|
|
await self.terminal.write(b"\x03")
|
|
await self.until(self.prompt)
|
|
await self.terminal.write(b"printf '\\nINTERRUPTED\\n'\r")
|
|
await self.until(b"\r\nINTERRUPTED\r\n")
|
|
|
|
async def test_explicit_end_terminates_the_shell(self):
|
|
await self.terminal.write(b"printf '\\nREADY\\n'\r")
|
|
await self.until(b"\r\nREADY\r\n")
|
|
await self.terminal.end()
|
|
async with asyncio.timeout(5):
|
|
while (message := await self.terminal.read())["type"] != "exit":
|
|
pass
|
|
self.assertIsInstance(message["exit_code"], int)
|
|
self.assertGreater(self.scope.kills, 0)
|
|
|
|
|
|
class TerminalCompatibilityTests(unittest.TestCase):
|
|
def test_old_agent_rejection_explains_how_to_upgrade(self):
|
|
left, right = socket.socketpair()
|
|
terminal = GuestTerminal(left)
|
|
self.addCleanup(terminal.close)
|
|
|
|
def old_agent():
|
|
with right:
|
|
agent.receive_frame(right, 4096, time.monotonic() + 2)
|
|
agent.send_frame(right, {"error": "invalid execution request fields", "outcome_unknown": False},
|
|
4096, time.monotonic() + 2)
|
|
|
|
thread = threading.Thread(target=old_agent)
|
|
thread.start()
|
|
try:
|
|
with self.assertRaisesRegex(GuestAgentUnavailableError, "Restart this VM to apply automatic guest-agent setup"):
|
|
terminal.start(24, 80)
|
|
finally:
|
|
thread.join(3)
|
|
|
|
|
|
class TerminalLifecycleTests(LifecycleTestCase):
|
|
def test_terminal_connects_under_lock_and_opens_without_holding_the_lock(self):
|
|
vm = self.create_vm(name="worker", ports=())
|
|
terminal = Mock()
|
|
|
|
def connect(path, vm_id, username, pid):
|
|
self.assertTrue(self.store.operation_active)
|
|
self.assertEqual((path, vm_id, username, pid), (self.settings.guest_socket_path(vm.id), vm.id, vm.username, vm.pid))
|
|
return terminal
|
|
|
|
def start(rows, cols):
|
|
self.assertFalse(self.store.operation_active)
|
|
self.assertEqual((rows, cols), (40, 120))
|
|
|
|
self.service._terminal_connector = Mock(side_effect=connect)
|
|
terminal.start.side_effect = start
|
|
self.assertEqual(self.service.open_terminal(vm.id, rows=40, cols=120), (terminal, vm.username))
|
|
terminal.close.assert_not_called()
|
|
|
|
def test_aliases_and_stopped_vms_cannot_create_a_terminal(self):
|
|
vm = self.create_vm(name="worker")
|
|
self.service._terminal_connector = Mock()
|
|
for alias in (vm.name, vm.guest_ip):
|
|
with self.assertRaises(NotFoundError):
|
|
self.service.open_terminal(alias)
|
|
self.service.stop(vm.id)
|
|
with self.assertRaises(ConflictError):
|
|
self.service.open_terminal(vm.id)
|
|
self.service._terminal_connector.assert_not_called()
|