mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-29 03:18:14 +08:00
198 lines
9.2 KiB
Python
198 lines
9.2 KiB
Python
"""用真实本地 Python 子进程验证持久 cell、IPC、输出与回收边界。"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import signal
|
||
import stat
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
import psutil
|
||
import pytest
|
||
import pytest_asyncio
|
||
|
||
from app.adapters.system.code.kernel import PythonKernel
|
||
from app.adapters.system.code.output import SPILL_BYTES, STDOUT_BYTES, OutputBuffer, project_stdout
|
||
from app.adapters.system.code.process import child_environment
|
||
from app.runtime.tasks import get_task_registry
|
||
|
||
|
||
async def _dispatch(name, arguments):
|
||
"""只返回调用参数，不触发任何真实业务或网络。"""
|
||
return {'name': name, 'value': arguments['value']}
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def kernel(tmp_path):
|
||
"""每个测试创建并可靠关闭真实内核，失败也不能把进程留给下一用例。"""
|
||
instance = PythonKernel(tmp_path / 'kernel', ('moviepilot_api',), blocking=asyncio.to_thread)
|
||
await instance.start()
|
||
try:
|
||
yield instance
|
||
finally:
|
||
await instance.close()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_python_state_imports_exceptions_and_fresh_rpc_context(kernel):
|
||
"""变量和导入跨 cell 保留，普通异常保留已赋值状态，每次 RPC 使用当前回调。"""
|
||
result = await kernel.execute(
|
||
'import math\nfrom moviepilot_tools import moviepilot_api\n'
|
||
'total = math.factorial(5)\nprint(moviepilot_api(value=total))', dispatch=_dispatch, timeout=5,
|
||
)
|
||
assert result['status'] == 'ok'
|
||
assert "'value': 120" in result['stdout']
|
||
failed = await kernel.execute('total += 1\nraise ValueError("cell failure")', dispatch=_dispatch, timeout=5)
|
||
assert failed['status'] == 'error' and 'ValueError: cell failure' in failed['traceback']
|
||
|
||
async def current(name, arguments):
|
||
"""替换回调即可改变当前调用身份，不重启 Python 或复用旧连接上下文。"""
|
||
return {'current': True, **await _dispatch(name, arguments)}
|
||
|
||
recovered = await kernel.execute('print(moviepilot_api(value=total))', dispatch=current, timeout=5)
|
||
assert recovered['execution_count'] == 3
|
||
assert "'current': True" in recovered['stdout'] and "'value': 121" in recovered['stdout']
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_multithreaded_rpc_does_not_mix_replies(kernel):
|
||
"""同一持久连接上的并发 Python 线程仍各自收到匹配的调用结果。"""
|
||
result = await kernel.execute(
|
||
'from concurrent.futures import ThreadPoolExecutor\n'
|
||
'from moviepilot_tools import moviepilot_api\n'
|
||
'with ThreadPoolExecutor(8) as pool:\n'
|
||
' print(list(pool.map(lambda n: moviepilot_api(value=n)["value"], range(30))))',
|
||
dispatch=_dispatch, timeout=5,
|
||
)
|
||
assert result['status'] == 'ok'
|
||
assert json.loads(result['stdout']) == list(range(30))
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fd_output_cannot_corrupt_frames_and_large_streams_are_drained(kernel):
|
||
"""直接 fd 输出与伪造帧文本不影响随机协议标记，大 stderr 不阻塞进程。"""
|
||
result = await kernel.execute(
|
||
'import os\nos.write(1, b"\\nFAKE_MARKER 999\\nraw output\\n")\n'
|
||
'os.write(2, b"E" * 100000)\nprint("cell result")', dispatch=_dispatch, timeout=5,
|
||
)
|
||
assert result['stdout'] == 'cell result\n'
|
||
assert 'FAKE_MARKER' in result['raw_stdout']
|
||
assert result['raw_stderr_bytes'] == 100000 and len(result['raw_stderr']) == 10000
|
||
next_result = await kernel.execute('print("next")', dispatch=_dispatch, timeout=5)
|
||
assert next_result['raw_stdout'] == next_result['raw_stderr'] == ''
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_token_allowlist_and_expired_cell_are_checked(kernel):
|
||
"""请求必须同时通过令牌、工具名和活动 cell 检查；旧派发引用不能跨 cell 使用。"""
|
||
good = {'token': kernel._token, 'tool': 'moviepilot_api', 'args': {'value': 1}}
|
||
assert 'No active' in (await kernel._request(json.dumps(good).encode(), None))['error']
|
||
kernel._dispatch = _dispatch
|
||
assert 'Unauthorized' in (await kernel._request(b'{"token":"wrong"}', _dispatch))['error']
|
||
assert 'Malformed' in (await kernel._request(b'not json', _dispatch))['error']
|
||
assert 'unavailable' in (await kernel._request(json.dumps({**good, 'tool': 'write_file'}).encode(), _dispatch))['error']
|
||
assert (await kernel._request(json.dumps(good).encode(), _dispatch))['value'] == 1
|
||
kernel._dispatch = None
|
||
assert 'No active' in (await kernel._request(json.dumps(good).encode(), _dispatch))['error']
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize('cancel', [False, True])
|
||
async def test_interrupted_cell_destroys_process_and_tasks(kernel, cancel):
|
||
"""超时及宿主取消都等到进程结束并释放端点，不允许继续使用不完整的状态。"""
|
||
task = get_task_registry().create(
|
||
kernel.execute('import time\ntime.sleep(60)', dispatch=_dispatch, timeout=0.1 if not cancel else 30),
|
||
owner='test.code.cell',
|
||
)
|
||
if cancel:
|
||
await asyncio.sleep(0.05)
|
||
task.cancel()
|
||
with pytest.raises(asyncio.CancelledError if cancel else TimeoutError):
|
||
await task
|
||
assert kernel.closed and kernel.process.returncode is not None
|
||
assert not kernel.directory.exists()
|
||
assert not [record for record in get_task_registry().records if record.owner.startswith('agent.code.')]
|
||
with pytest.raises(RuntimeError, match='unavailable'):
|
||
await kernel.execute('print("must not run")', dispatch=_dispatch, timeout=5)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_cancellation_during_file_creation_cleans_late_work(tmp_path, monkeypatch):
|
||
"""线程创建不能被取消；宿主必须等创建真正结束，再回收迟到文件。"""
|
||
instance = PythonKernel(tmp_path / 'late', (), blocking=asyncio.to_thread)
|
||
entered, release = threading.Event(), threading.Event()
|
||
original = instance._files
|
||
|
||
def delayed():
|
||
"""在创建线程持有一个可由测试控制的有限等待点。"""
|
||
entered.set()
|
||
assert release.wait(5)
|
||
original()
|
||
|
||
monkeypatch.setattr(instance, '_files', delayed)
|
||
task = get_task_registry().create(instance.start(), owner='test.code.start')
|
||
try:
|
||
assert await asyncio.to_thread(entered.wait, 5)
|
||
task.cancel()
|
||
release.set()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await task
|
||
assert not instance.directory.exists()
|
||
assert instance.process is None and instance.server is None
|
||
finally:
|
||
release.set()
|
||
await instance.close()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_private_files_and_parent_death_signal(kernel):
|
||
"""私有文件不可被其他系统用户读取，父死亡管道在正在执行时也能终止内核。"""
|
||
if os.name == 'nt':
|
||
pytest.skip('POSIX 文件权限和死亡管道；Windows 使用进程句柄')
|
||
assert stat.S_IMODE(kernel.directory.stat().st_mode) == 0o700
|
||
assert stat.S_IMODE((kernel.directory / 'moviepilot_tools.py').stat().st_mode) == 0o600
|
||
assert stat.S_IMODE(Path(kernel.endpoint).stat().st_mode) == 0o600
|
||
kernel.process.stdin.write(b'{"id":"death", "code":"import time; time.sleep(60)"}\n')
|
||
await kernel.process.stdin.drain()
|
||
kernel._watch.close()
|
||
await asyncio.wait_for(kernel.process.wait(), 5)
|
||
assert kernel.process.returncode == -signal.SIGKILL
|
||
|
||
|
||
def test_environment_drops_secrets_before_safe_prefixes():
|
||
"""敏感名称优先于安全前缀；内核不能继承宿主导入路径或认证令牌。"""
|
||
environment = child_environment({'PATH': '/bin', 'HOME': '/home/example', 'HOME_TOKEN': 'secret',
|
||
'LANG': 'en_US.UTF-8', 'OPENAI_API_KEY': 'secret', 'PYTHONPATH': '/host'})
|
||
assert environment['PATH'] == '/bin' and environment['HOME'] == '/home/example'
|
||
assert not {'HOME_TOKEN', 'OPENAI_API_KEY', 'PYTHONPATH'} & environment.keys()
|
||
assert environment['PYTHONUTF8'] == '1'
|
||
|
||
|
||
def test_bounded_output_preserves_readable_spill(tmp_path):
|
||
"""预览按字节预算保留首尾，正文可续读且超过5MB时明确报告上限。"""
|
||
text = 'A' * STDOUT_BYTES + '中' * SPILL_BYTES + 'THE_END'
|
||
preview, metadata = project_stdout(text, tmp_path)
|
||
assert preview.startswith('AAAA') and preview.endswith('THE_END')
|
||
assert metadata['stdout_truncated'] and metadata['stdout_spill_capped']
|
||
assert Path(metadata['stdout_spill_path']).stat().st_size == SPILL_BYTES
|
||
assert metadata['stdout_bytes_total'] == len(text.encode())
|
||
buffer = OutputBuffer(10)
|
||
buffer.append(b'x' * 100)
|
||
assert buffer.drain() == ('x' * 10, 100)
|
||
assert buffer.drain() == ('', 0)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_parent_death_also_retires_inherited_descendants(kernel):
|
||
"""父进程死亡不能留下继承输出管道的孙进程，使宿主一直等不到 EOF。"""
|
||
if os.name == 'nt':
|
||
pytest.skip('真实 POSIX 进程组；Windows 由 taskkill 树回收')
|
||
result = await kernel.execute('import subprocess, sys\n'
|
||
'child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"])\n'
|
||
'print(child.pid)', dispatch=_dispatch, timeout=5)
|
||
child_pid = int(result['stdout'])
|
||
kernel._watch.close()
|
||
await asyncio.wait_for(kernel.process.wait(), 5)
|
||
assert not psutil.pid_exists(child_pid) or psutil.Process(child_pid).status() == psutil.STATUS_ZOMBIE
|