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 'REAL-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"sleep 20\r") await asyncio.sleep(0.1) await self.terminal.write(b"\x03") await self.until(self.prompt) await self.terminal.write(b"printf 'INTERRUPTED\\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 'READY\\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, "Install / update terminal support"): 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()