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()