184 lines
10 KiB
Python
184 lines
10 KiB
Python
from __future__ import annotations
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
from core.sandbox.capacity import ExecCapacity
|
|
from core.sandbox.pool import SandboxPool
|
|
from core.sandbox.package_scans import persist_scan
|
|
from deploy.sandbox.package_scan import scan_packages
|
|
from core import procs
|
|
from web.admin import _sandbox_package_stats
|
|
|
|
|
|
class CapacityTests(unittest.TestCase):
|
|
def test_env_overrides_capacity_config(self):
|
|
with tempfile.TemporaryDirectory() as td, patch.dict(os.environ, {
|
|
"ZCBOT_MAX_ACTIVE_EXECS": "5", "ZCBOT_MAX_BACKGROUND_EXECS": "3",
|
|
"ZCBOT_MAX_ACTIVE_EXECS_PER_USER": "1", "ZCBOT_MIN_MEM_AVAILABLE": "2g",
|
|
}, clear=False):
|
|
cap = ExecCapacity(Path(td), {"max_active_execs": 9, "max_background_execs": 8,
|
|
"max_active_execs_per_user": 7, "min_mem_available": "4g"})
|
|
self.assertEqual((cap.max_active, cap.max_background, cap.max_per_user, cap.min_mem_available),
|
|
(5, 3, 1, 2 * 1024**3))
|
|
|
|
def test_config_cannot_raise_hard_limits(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
cap = ExecCapacity(Path(td), {"max_active_execs": 99, "max_background_execs": 99,
|
|
"max_active_execs_per_user": 99})
|
|
self.assertEqual((cap.max_active, cap.max_background, cap.max_per_user), (6, 4, 2))
|
|
|
|
def test_global_per_user_and_background_limits(self):
|
|
with tempfile.TemporaryDirectory() as td, patch("core.sandbox.capacity.mem_available_bytes", return_value=10**12):
|
|
cap = ExecCapacity(Path(td), {"max_active_execs": 3, "max_background_execs": 1, "max_active_execs_per_user": 2})
|
|
a = cap.try_acquire("u1", "background", lease_id="a")
|
|
b = cap.try_acquire("u1", "foreground", lease_id="b")
|
|
self.assertIsNotNone(a); self.assertIsNotNone(b)
|
|
self.assertIsNone(cap.try_acquire("u1", "foreground", lease_id="c"))
|
|
self.assertIsNone(cap.try_acquire("u2", "background", lease_id="d"))
|
|
self.assertIsNotNone(cap.try_acquire("u2", "foreground", lease_id="e"))
|
|
self.assertIsNone(cap.try_acquire("u3", "foreground", lease_id="f"))
|
|
|
|
def test_cross_instance_state_and_cancel_queue(self):
|
|
with tempfile.TemporaryDirectory() as td, patch("core.sandbox.capacity.mem_available_bytes", return_value=10**12):
|
|
one = ExecCapacity(Path(td), {"max_active_execs": 1})
|
|
two = ExecCapacity(Path(td), {"max_active_execs": 1})
|
|
lease = one.try_acquire("u1", "foreground", lease_id="held")
|
|
cancelled = threading.Event()
|
|
result = []
|
|
t = threading.Thread(target=lambda: result.append(two.acquire_foreground("u2", cancelled.is_set)))
|
|
t.start(); time.sleep(.15); cancelled.set(); t.join(2)
|
|
self.assertEqual(result, [None])
|
|
self.assertEqual(two.snapshot()["foreground_queued"], 0)
|
|
one.release(lease)
|
|
|
|
def test_memory_pressure_pauses_new_admission(self):
|
|
with tempfile.TemporaryDirectory() as td, patch("core.sandbox.capacity.mem_available_bytes", return_value=100):
|
|
cap = ExecCapacity(Path(td), {"min_mem_available": "1g"})
|
|
self.assertIsNone(cap.try_acquire("u", "foreground"))
|
|
self.assertTrue(cap.snapshot()["memory_paused"])
|
|
|
|
|
|
class PoolTests(unittest.TestCase):
|
|
def test_defaults_and_docker_tmpfs(self):
|
|
with tempfile.TemporaryDirectory() as td, patch("core.sandbox.pool.subprocess.run") as run:
|
|
run.return_value = MagicMock(returncode=0, stdout="")
|
|
pool = SandboxPool(Path(td) / "users")
|
|
self.assertEqual((pool.memory, pool.cpus, pool.idle_ttl, pool.tmp_size), ("4g", "2.0", 600, "1g"))
|
|
pool._docker_run(uuid4(), "box")
|
|
argv = run.call_args_list[-1].args[0]
|
|
self.assertIn("/tmp:exec,size=1g,mode=1777", argv)
|
|
|
|
def test_active_exec_not_reaped(self):
|
|
with tempfile.TemporaryDirectory() as td, patch("core.sandbox.pool._container_running", return_value=True), patch("core.sandbox.pool.subprocess.run") as run:
|
|
pool = SandboxPool(Path(td) / "users", idle_ttl=1)
|
|
uid = uuid4(); pool._last_active[uid] = int(time.time()) - 10
|
|
pool.exec_started(uid)
|
|
self.assertEqual(pool.reap_idle(), [])
|
|
run.assert_not_called()
|
|
|
|
def test_scan_failure_does_not_block_idle_removal(self):
|
|
with tempfile.TemporaryDirectory() as td, \
|
|
patch("core.sandbox.package_scans.scan_and_persist", side_effect=RuntimeError("db down")), \
|
|
patch("core.sandbox.pool.subprocess.run", return_value=MagicMock(returncode=0)):
|
|
pool = SandboxPool(Path(td) / "users", idle_ttl=1)
|
|
uid = uuid4(); pool._last_active[uid] = int(time.time()) - 10
|
|
self.assertEqual(pool.reap_idle(), [f"zcbot-sandbox-{uid}"])
|
|
|
|
def test_finalizer_helper_swallows_scanner_error(self):
|
|
with tempfile.TemporaryDirectory() as td, patch("core.sandbox.package_scans.scan_and_persist", side_effect=RuntimeError("db down")):
|
|
pool = SandboxPool(Path(td) / "users")
|
|
pool._scan_before_remove("box", uuid4(), "foreground")
|
|
|
|
|
|
class PackageScannerTests(unittest.TestCase):
|
|
def test_empty_scan_does_not_touch_database(self):
|
|
persist_scan(container_session_id="empty", user_id=uuid4(), execution_kind="foreground",
|
|
result={"packages": [], "image_digest": "sha256:x", "python_version": "3.12", "total_installed_bytes": 0})
|
|
|
|
def test_added_override_bytes_and_symlink_escape(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
root = Path(td) / "lib"; site = root / "python3.12" / "site-packages"; site.mkdir(parents=True)
|
|
pkg = site / "demo.py"; pkg.write_text("abc", encoding="utf-8")
|
|
dist = site / "Demo-2.dist-info"; dist.mkdir()
|
|
(dist / "METADATA").write_text("Name: Demo\nVersion: 2\n", encoding="utf-8")
|
|
(dist / "REQUESTED").write_text("", encoding="utf-8")
|
|
(dist / "RECORD").write_text("demo.py,,3\n", encoding="utf-8")
|
|
rows = scan_packages(root, {"demo": "1"})
|
|
self.assertEqual(rows[0]["change"], "override")
|
|
self.assertEqual(rows[0]["installed_bytes"], 3)
|
|
self.assertTrue(rows[0]["direct_requested"])
|
|
self.assertEqual(scan_packages(root, {"demo": "2"})[0]["change"], "reinstall")
|
|
self.assertEqual(scan_packages(root, {})[0]["change"], "added")
|
|
|
|
def test_metadata_symlink_is_ignored(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
root = Path(td) / "lib"; site = root / "python3.12" / "site-packages"; site.mkdir(parents=True)
|
|
outside = Path(td) / "outside"; outside.write_text("Name: Bad\nVersion: 1\n")
|
|
dist = site / "Bad.dist-info"; dist.mkdir()
|
|
try:
|
|
(dist / "METADATA").symlink_to(outside)
|
|
except OSError:
|
|
self.skipTest("symlink unavailable")
|
|
self.assertEqual(scan_packages(root, {}), [])
|
|
|
|
|
|
class ProcQueueTests(unittest.TestCase):
|
|
def test_queued_status_and_cancel(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
d = Path(td) / "p"; d.mkdir()
|
|
meta = {"proc_id": "p1", "backend": "docker", "state": "queued"}
|
|
procs.write_meta(d, meta)
|
|
self.assertEqual(procs.status_of(meta, d)[0], "queued")
|
|
self.assertEqual(procs.kill_proc(meta, d), "已取消排队")
|
|
self.assertEqual(procs.status_of(procs.read_meta(d), d), ("finished", 137))
|
|
|
|
def test_start_queued_respects_capacity_then_runs(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
d = Path(td) / "p"; d.mkdir(); (d / "runner.sh").write_text("true")
|
|
uid = uuid4()
|
|
meta = {"proc_id": "p1", "task_id": "t1", "backend": "docker", "state": "queued",
|
|
"user_id": str(uid), "cwd": "/workspace/demo"}
|
|
procs.write_meta(d, meta)
|
|
pool = MagicMock(); pool.capacity.try_acquire.return_value = None
|
|
self.assertFalse(procs.start_queued_docker(meta, d, pool))
|
|
pool.capacity.try_acquire.return_value = "bg-p1"
|
|
pool.run_proc_container.return_value = "zcbot-proc-p1"
|
|
with patch("core.procs.subprocess.run", return_value=MagicMock(returncode=0, stderr="")):
|
|
self.assertTrue(procs.start_queued_docker(meta, d, pool))
|
|
saved = procs.read_meta(d)
|
|
self.assertEqual(saved["state"], "running")
|
|
self.assertEqual(saved["capacity_lease_id"], "bg-p1")
|
|
|
|
|
|
class AdminAggregationTests(unittest.TestCase):
|
|
def test_package_sessions_are_aggregated_by_name_version(self):
|
|
scans = [
|
|
SimpleNamespace(container_session_id="s1", user_id=uuid4(), execution_kind="foreground",
|
|
packages=[{"name":"Demo","version":"1","direct_requested":True,"installed_bytes":100,"change":"added","base_version":None}],
|
|
finished_at=datetime.now(timezone.utc)),
|
|
SimpleNamespace(container_session_id="s2", user_id=uuid4(), execution_kind="background",
|
|
packages=[{"name":"Demo","version":"1","direct_requested":False,"installed_bytes":300,"change":"added","base_version":None}],
|
|
finished_at=datetime.now(timezone.utc)),
|
|
]
|
|
session = MagicMock(); session.execute.return_value.scalars.return_value.all.return_value = scans
|
|
row = _sandbox_package_stats(session, None)["rows"][0]
|
|
self.assertEqual(row["session_count"], 2)
|
|
self.assertEqual(row["user_count"], 2)
|
|
self.assertEqual(row["foreground_sessions"], 1)
|
|
self.assertEqual(row["background_sessions"], 1)
|
|
self.assertEqual(row["direct_sessions"], 1)
|
|
self.assertEqual(row["average_installed_bytes"], 200)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|