327 lines
12 KiB
Python
327 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import json
|
|
import os
|
|
import socket
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
SDK_SRC = PROJECT_ROOT / "packages" / "jwebot-sdk" / "src"
|
|
KERNEL_SRC = PROJECT_ROOT / "packages" / "jwebot-kernel" / "src"
|
|
sys.path[:0] = [str(SDK_SRC), str(KERNEL_SRC)]
|
|
|
|
from jwebot_kernel import ( # noqa: E402
|
|
ExtensionHostClient,
|
|
HostClientError,
|
|
HostTimeout,
|
|
RuntimeTelemetry,
|
|
UnixExtensionHostClient,
|
|
)
|
|
|
|
|
|
class HostClientTests(unittest.TestCase):
|
|
def helper(self, source: str):
|
|
temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temporary.cleanup)
|
|
root = Path(temporary.name)
|
|
release = root / "release"
|
|
release.mkdir()
|
|
script = root / "host.py"
|
|
script.write_text(source, encoding="utf-8")
|
|
return root, release, script
|
|
|
|
def test_client_uses_minimal_environment_and_parses_response(self):
|
|
root, release, script = self.helper(
|
|
"\n".join(
|
|
(
|
|
"import json, os, sys",
|
|
"request = json.loads(sys.stdin.readline())",
|
|
"print(json.dumps({",
|
|
" 'status': 'ok',",
|
|
" 'event_id': request['event_id'],",
|
|
" 'secret_present': 'JWEBOT_TEST_SECRET' in os.environ,",
|
|
" 'dont_write_bytecode': sys.dont_write_bytecode,",
|
|
"}), flush=True)",
|
|
)
|
|
)
|
|
)
|
|
del root
|
|
os.environ["JWEBOT_TEST_SECRET"] = "must-not-cross"
|
|
self.addCleanup(os.environ.pop, "JWEBOT_TEST_SECRET", None)
|
|
response = ExtensionHostClient(
|
|
command=(sys.executable, str(script)),
|
|
).invoke(release, {"event_id": "event-1"})
|
|
self.assertEqual(response["event_id"], "event-1")
|
|
self.assertFalse(response["secret_present"])
|
|
self.assertTrue(response["dont_write_bytecode"])
|
|
|
|
def test_client_rejects_invalid_json(self):
|
|
_, release, script = self.helper("print('not-json')\n")
|
|
with self.assertRaisesRegex(HostClientError, "invalid JSON"):
|
|
ExtensionHostClient(
|
|
command=(sys.executable, str(script)),
|
|
).invoke(release, {"event_id": "event-1"})
|
|
|
|
def test_client_rejects_host_error_status(self):
|
|
_, release, script = self.helper(
|
|
"print('{\"status\":\"error\",\"error\":\"ExtensionFailed\"}')\n"
|
|
)
|
|
with self.assertRaisesRegex(HostClientError, "returned an error"):
|
|
ExtensionHostClient(
|
|
command=(sys.executable, str(script)),
|
|
).invoke(release, {"event_id": "event-1"})
|
|
|
|
def test_client_enforces_request_limit(self):
|
|
_, release, script = self.helper("print('{\"status\":\"ok\"}')\n")
|
|
with self.assertRaisesRegex(HostClientError, "size limit"):
|
|
ExtensionHostClient(
|
|
command=(sys.executable, str(script)),
|
|
).invoke(release, {"value": "x" * (256 * 1024)})
|
|
|
|
def test_timeout_kills_the_host_process_group(self):
|
|
root, release, script = self.helper(
|
|
"\n".join(
|
|
(
|
|
"import subprocess, sys, time",
|
|
"marker = sys.argv[1]",
|
|
"subprocess.Popen([",
|
|
" sys.executable, '-c',",
|
|
" 'import pathlib,time,sys; time.sleep(0.3); '",
|
|
" 'pathlib.Path(sys.argv[1]).write_text(\"escaped\")',",
|
|
" marker,",
|
|
"])",
|
|
"time.sleep(10)",
|
|
)
|
|
)
|
|
)
|
|
marker = root / "escaped.txt"
|
|
client = ExtensionHostClient(
|
|
command=(sys.executable, str(script), str(marker)),
|
|
timeout_seconds=0.05,
|
|
)
|
|
with self.assertRaises(HostTimeout):
|
|
client.invoke(release, {"event_id": "event-1"})
|
|
time.sleep(0.5)
|
|
self.assertFalse(marker.exists())
|
|
|
|
def test_client_services_one_bounded_broker_request(self):
|
|
_, release, script = self.helper(
|
|
"\n".join(
|
|
(
|
|
"import json, sys",
|
|
"request = json.loads(sys.stdin.readline())",
|
|
"print(json.dumps({",
|
|
" 'status': 'service_request',",
|
|
" 'call_id': '1',",
|
|
" 'profile': 'coin',",
|
|
" 'operation': 'execute',",
|
|
" 'payload': {'command': '/翻硬币'},",
|
|
"}), flush=True)",
|
|
"service = json.loads(sys.stdin.readline())",
|
|
"print(json.dumps({",
|
|
" 'status': 'ok',",
|
|
" 'plugin_id': 'coin',",
|
|
" 'request_op': request['op'],",
|
|
" 'text': service['result']['text'],",
|
|
"}), flush=True)",
|
|
)
|
|
)
|
|
)
|
|
calls = []
|
|
observed = []
|
|
|
|
def service(profile, operation, payload):
|
|
calls.append((profile, operation, payload))
|
|
return {"text": "硬币落地：正面喵"}
|
|
|
|
response = ExtensionHostClient(
|
|
command=(sys.executable, str(script)),
|
|
plugin_id="coin",
|
|
telemetry=RuntimeTelemetry(
|
|
component="host-client",
|
|
observer=observed.append,
|
|
logger=mock.Mock(),
|
|
),
|
|
).invoke(
|
|
release,
|
|
{"op": "invoke_command"},
|
|
service_handler=service,
|
|
)
|
|
|
|
self.assertEqual(response["text"], "硬币落地：正面喵")
|
|
self.assertEqual(
|
|
calls,
|
|
[
|
|
(
|
|
"coin",
|
|
"execute",
|
|
{"command": "/翻硬币"},
|
|
)
|
|
],
|
|
)
|
|
names = [value["name"] for value in observed]
|
|
self.assertIn("host.process.spawn", names)
|
|
self.assertIn("host.service.request", names)
|
|
self.assertIn("host.service.complete", names)
|
|
self.assertIn("host.invoke", names)
|
|
|
|
def test_service_handler_cannot_exceed_host_deadline(self):
|
|
_, release, script = self.helper(
|
|
"\n".join(
|
|
(
|
|
"import json, sys",
|
|
"json.loads(sys.stdin.readline())",
|
|
"print(json.dumps({",
|
|
" 'status': 'service_request',",
|
|
" 'call_id': '1',",
|
|
" 'profile': 'coin',",
|
|
" 'operation': 'execute',",
|
|
" 'payload': {},",
|
|
"}), flush=True)",
|
|
"sys.stdin.readline()",
|
|
)
|
|
)
|
|
)
|
|
|
|
def slow_service(profile, operation, payload):
|
|
del profile, operation, payload
|
|
time.sleep(0.65)
|
|
return {"text": "late"}
|
|
|
|
with self.assertRaisesRegex(HostTimeout, "service exceeded"):
|
|
ExtensionHostClient(
|
|
command=(sys.executable, str(script)),
|
|
timeout_seconds=0.5,
|
|
).invoke(
|
|
release,
|
|
{"op": "invoke_command"},
|
|
service_handler=slow_service,
|
|
)
|
|
|
|
|
|
class UnixHostClientTests(unittest.TestCase):
|
|
def server(self, response: bytes, *, delay: float = 0.0):
|
|
temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temporary.cleanup)
|
|
path = Path(temporary.name) / "host.sock"
|
|
listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
listener.bind(str(path))
|
|
listener.listen(1)
|
|
received: list[bytes] = []
|
|
|
|
def serve():
|
|
try:
|
|
connection, _ = listener.accept()
|
|
with connection:
|
|
stream = connection.makefile("rwb", buffering=0)
|
|
received.append(stream.readline(256 * 1024 + 1))
|
|
if delay:
|
|
time.sleep(delay)
|
|
try:
|
|
stream.write(response)
|
|
except OSError:
|
|
pass
|
|
stream.close()
|
|
finally:
|
|
listener.close()
|
|
|
|
thread = threading.Thread(target=serve, daemon=True)
|
|
thread.start()
|
|
self.addCleanup(thread.join, 1)
|
|
return path, received
|
|
|
|
def test_unix_client_round_trip(self):
|
|
path, received = self.server(
|
|
b'{"status":"ok","plugin_id":"hello-command"}\n'
|
|
)
|
|
response = UnixExtensionHostClient(path).invoke(
|
|
{"op": "invoke_command", "command": "/hello"}
|
|
)
|
|
self.assertEqual(response["plugin_id"], "hello-command")
|
|
self.assertEqual(
|
|
json.loads(received[0]),
|
|
{"command": "/hello", "op": "invoke_command"},
|
|
)
|
|
|
|
def test_unix_client_timeout_is_bounded(self):
|
|
path, _ = self.server(
|
|
b'{"status":"ok"}\n',
|
|
delay=0.3,
|
|
)
|
|
with self.assertRaises(HostTimeout):
|
|
UnixExtensionHostClient(
|
|
path,
|
|
timeout_seconds=0.05,
|
|
).invoke({"op": "invoke_command"})
|
|
|
|
def test_unix_client_rejects_error_response(self):
|
|
path, _ = self.server(
|
|
b'{"status":"error","error":"ExtensionExecutionError"}\n'
|
|
)
|
|
with self.assertRaisesRegex(HostClientError, "returned an error"):
|
|
UnixExtensionHostClient(path).invoke({"op": "invoke_command"})
|
|
|
|
def test_unix_client_fails_closed_when_socket_is_missing(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "missing.sock"
|
|
with self.assertRaisesRegex(HostClientError, "unavailable"):
|
|
UnixExtensionHostClient(path).invoke(
|
|
{"op": "invoke_command"}
|
|
)
|
|
|
|
def test_unix_client_services_broker_request(self):
|
|
temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temporary.cleanup)
|
|
path = Path(temporary.name) / "host.sock"
|
|
listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
listener.bind(str(path))
|
|
listener.listen(1)
|
|
received = []
|
|
|
|
def serve():
|
|
try:
|
|
connection, _ = listener.accept()
|
|
with connection:
|
|
stream = connection.makefile("rwb", buffering=0)
|
|
received.append(json.loads(stream.readline()))
|
|
stream.write(
|
|
b'{"status":"service_request","call_id":"1",'
|
|
b'"profile":"coin","operation":"execute",'
|
|
b'"payload":{"command":"/\\u7ffb\\u786c\\u5e01"}}\n'
|
|
)
|
|
received.append(json.loads(stream.readline()))
|
|
stream.write(
|
|
b'{"status":"ok","plugin_id":"coin"}\n'
|
|
)
|
|
stream.close()
|
|
finally:
|
|
listener.close()
|
|
|
|
thread = threading.Thread(target=serve, daemon=True)
|
|
thread.start()
|
|
self.addCleanup(thread.join, 1)
|
|
calls = []
|
|
|
|
response = UnixExtensionHostClient(path).invoke(
|
|
{"op": "invoke_command", "command": "/翻硬币"},
|
|
service_handler=lambda profile, operation, payload: (
|
|
calls.append((profile, operation, payload))
|
|
or {"text": "ok"}
|
|
),
|
|
)
|
|
|
|
self.assertEqual(response["plugin_id"], "coin")
|
|
self.assertEqual(received[1]["op"], "service_result")
|
|
self.assertEqual(calls[0][0:2], ("coin", "execute"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|