"""Exercise the production collector with disposable local process groups.""" import importlib.util import io import json import os from pathlib import Path import subprocess import sys import tempfile import threading import time import unittest import uuid from unittest.mock import patch DIRECTORY = Path(__file__).resolve().parents[3] / 'deploy' / 'computer' sys.path.insert(0, str(DIRECTORY)) spec = importlib.util.spec_from_file_location('compose_bridge', DIRECTORY / 'compose_bridge.py') bridge = importlib.util.module_from_spec(spec) spec.loader.exec_module(bridge) def start_local(argv, cwd, user='cua'): return subprocess.Popen(argv, cwd=cwd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, start_new_session=True) def request(path, body): data = json.dumps({'desktopId': 'test', 'token': 'token', 'instance': bridge.INSTANCE, **body}).encode() handler = object.__new__(bridge.Handler) handler.path = path handler.headers = {'Content-Length': str(len(data))} handler.rfile, handler.wfile = io.BytesIO(data), io.BytesIO() statuses = [] handler.send_response = statuses.append handler.send_header = lambda *args: None handler.end_headers = lambda: None handler.do_POST() return statuses[0], json.loads(handler.wfile.getvalue()) class CommandChecks(unittest.TestCase): def test_cancel_stops_only_its_process_group_and_keeps_peer_and_lease(self): with tempfile.TemporaryDirectory() as directory: ready, late, peer = [Path(directory) / name for name in ('ready', 'late', 'peer')] child = "import pathlib,time;time.sleep(.6);pathlib.Path(%r).write_text('late')" % str(late) script = "import pathlib,subprocess,sys,time;subprocess.Popen([sys.executable,'-c',%r]);pathlib.Path(%r).touch();time.sleep(30)" % (child, str(ready)) peer_script = "import pathlib,time;time.sleep(.3);pathlib.Path(%r).write_text('ok')" % str(peer) lease = {'token': 'token', 'expires': time.monotonic() + 90, 'active': 0} key, results, started = str(uuid.uuid4()), {}, [] def launch(argv, cwd, user='cua'): process = start_local(argv, directory) started.append(process) return process with patch.object(bridge, 'DESKTOP_ID', 'test'), patch.object(bridge, 'LEASE', lease), \ patch.object(bridge, 'start', side_effect=launch), patch.object(bridge, 'reset_desktop') as reset: threads = [threading.Thread(target=lambda: results.update(cancelled=request('/execute', {'operationId': key, 'argv': [sys.executable, '-c', script]}))), threading.Thread(target=lambda: results.update(peer=request('/execute', {'operationId': str(uuid.uuid4()), 'argv': [sys.executable, '-c', peer_script]})))] try: for thread in threads: thread.start() deadline = time.monotonic() + 3 while not ready.exists() and time.monotonic() < deadline: time.sleep(.01) self.assertTrue(ready.exists()) self.assertEqual(request('/cancel', {'operationId': key}), (200, {'ok': True, 'stopped': True, 'operationId': key})) for thread in threads: thread.join(3) self.assertFalse(thread.is_alive()) self.assertEqual(results['cancelled'][0], 409) self.assertEqual(results['peer'][0], 200) self.assertEqual(peer.read_text(), 'ok') time.sleep(.5) self.assertFalse(late.exists(), 'cancel left a child running') self.assertEqual(lease['active'], 0) self.assertIs(bridge.LEASE, lease) reset.assert_not_called() finally: for process in started: bridge.stop_command(process) for thread in threads: thread.join(3) def test_cancel_before_execute_prevents_late_launch_and_requires_the_current_lease(self): lease = {'token': 'token', 'expires': time.monotonic() + 90, 'active': 0} key = str(uuid.uuid4()) with patch.object(bridge, 'DESKTOP_ID', 'test'), patch.object(bridge, 'LEASE', lease), patch.object(bridge, 'start') as launch: self.assertEqual(request('/cancel', {'operationId': key, 'token': 'wrong'})[0], 409) self.assertEqual(request('/cancel', {'operationId': key})[0], 200) self.assertEqual(request('/execute', {'operationId': key, 'argv': ['never-start']})[0], 409) self.assertEqual(request('/cancel', {'operationId': key})[0], 200) launch.assert_not_called() def test_timeout_does_not_restart_another_holders_desktop(self): with tempfile.TemporaryDirectory() as directory: marker = Path(directory) / 'should-not-exist' script = "import pathlib,time; time.sleep(.4); pathlib.Path(%r).write_text('late')" % str(marker) command = "import subprocess,sys,time; subprocess.Popen([sys.executable,'-c',%r]); time.sleep(30)" % script # Keep one peer active; exercise the actual HTTP handler error path. lease = {'token': 'token', 'expires': time.monotonic() + 90, 'active': 1} body = json.dumps({'desktopId': 'test', 'token': 'token', 'instance': bridge.INSTANCE, 'argv': [sys.executable, '-c', command], 'timeout': 1000}).encode() handler = object.__new__(bridge.Handler) handler.path = '/execute' handler.headers = {'Content-Length': str(len(body))} handler.rfile = io.BytesIO(body) handler.wfile = io.BytesIO() statuses = [] handler.send_response = statuses.append handler.send_header = lambda *args: None handler.end_headers = lambda: None # Make the timeout short enough to catch a surviving child before it writes. started = [] def launch(argv, cwd, user='cua'): process = start_local(argv, directory) started.append(process) return process original = bridge.collect with patch.object(bridge, 'DESKTOP_ID', 'test'), patch.object(bridge, 'LEASE', lease), \ patch.object(bridge, 'start', side_effect=launch), \ patch.object(bridge, 'collect', side_effect=lambda p, data, timeout: original(p, data, .15)), \ patch.object(bridge, 'reset_desktop') as reset: try: handler.do_POST() self.assertEqual(statuses, [409]) reset.assert_not_called() self.assertIs(bridge.LEASE, lease) self.assertEqual(lease['active'], 1) self.assertIsNotNone(started[0].poll()) time.sleep(.5) self.assertFalse(marker.exists(), 'timed out command left a child process running') finally: for process in started: try: os.killpg(process.pid, 9) except ProcessLookupError: pass process.wait() def test_output_cap_is_enforced_while_the_process_is_still_writing(self): with patch.object(bridge, 'start', side_effect=start_local), patch.object(bridge, 'reset_desktop') as reset: before = time.monotonic() with self.assertRaisesRegex(ValueError, '输出超过限制'): bridge.run([sys.executable, '-c', "import os,time; os.write(1,b'x'*(9*1024*1024)); time.sleep(10)"], cwd='.', timeout=3) self.assertLess(time.monotonic() - before, 2) reset.assert_not_called() def test_large_input_and_both_output_streams_make_progress(self): with patch.object(bridge, 'start', side_effect=start_local): script = "import sys; sys.stdout.write('o'*200000); sys.stdout.flush(); sys.stderr.write('e'*200000); sys.stderr.flush(); data=sys.stdin.buffer.read(); print(len(data))" result = bridge.run([sys.executable, '-c', script], cwd='.', data=b'i' * 1000000, timeout=3) self.assertEqual(result['exitCode'], 0) self.assertEqual(result['stdout'], 'o' * 200000 + '1000000\n') self.assertEqual(result['stderr'], 'e' * 200000) if __name__ == '__main__': result = unittest.TextTestRunner().run(unittest.defaultTestLoader.loadTestsFromTestCase(CommandChecks)) if not result.wasSuccessful(): raise SystemExit(1) print('COMPOSE_BRIDGE_CHECKS_OK')