113 lines
4.7 KiB
Python
113 lines
4.7 KiB
Python
"""Guards for the memory layout described in docs/memory.md.
|
|
|
|
* The server process must not import the math libraries (they live in the workers).
|
|
* Workers must not import the MCP/HTTP stack (spawn re-imports the main module).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import builtins
|
|
import io
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import textwrap
|
|
import typing
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from mcp_calculator import executor as executor_module
|
|
from mcp_calculator.config import load_settings
|
|
from mcp_calculator.executor import ComputationError, Executor
|
|
from mcp_calculator.ops import spec, stats
|
|
from tests.support import slow
|
|
|
|
MATH_LIBS = ("sympy", "numpy", "scipy", "pint", "mpmath")
|
|
SERVER_LIBS = ("mcp", "uvicorn", "starlette", "pydantic", "anyio")
|
|
|
|
|
|
def _run_python(code: str, modules: tuple[str, ...]) -> dict:
|
|
"""Run ``code`` as a script (MODULES is replaced by the module names); return its last JSON line."""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
script = Path(tmp) / "probe.py"
|
|
script.write_text(textwrap.dedent(code).replace("MODULES", repr(modules)))
|
|
out = subprocess.run([sys.executable, str(script)], capture_output=True, text=True, timeout=120, check=True)
|
|
return json.loads(out.stdout.strip().splitlines()[-1])
|
|
|
|
|
|
@slow
|
|
class ImportFootprintTest(unittest.TestCase):
|
|
def test_server_process_does_not_import_math_libraries(self):
|
|
loaded = _run_python(
|
|
"""
|
|
import json, sys
|
|
import mcp_calculator.__main__, mcp_calculator.server, mcp_calculator.http_app, mcp_calculator.executor
|
|
print(json.dumps({m: m in sys.modules for m in MODULES}))
|
|
""",
|
|
MATH_LIBS,
|
|
)
|
|
self.assertFalse(any(loaded.values()), loaded)
|
|
|
|
def test_workers_do_not_import_server_stack(self):
|
|
# Mirrors the console script: the main module imports mcp_calculator.__main__.
|
|
loaded = _run_python(
|
|
"""
|
|
import json, sys
|
|
from mcp_calculator.__main__ import main # noqa: F401
|
|
|
|
def probe():
|
|
import os
|
|
loaded = {m: m in sys.modules for m in MODULES}
|
|
loaded["single_threaded_blas"] = os.environ.get("OPENBLAS_NUM_THREADS") == "1"
|
|
return loaded
|
|
|
|
if __name__ == "__main__":
|
|
from mcp_calculator.config import load_settings
|
|
from mcp_calculator.executor import Executor
|
|
|
|
executor = Executor(load_settings([], env={"MCP_CALC_WORKERS": "1"}))
|
|
executor.start()
|
|
result = executor._pool.schedule(probe, timeout=60).result()
|
|
executor.stop()
|
|
print(json.dumps(result))
|
|
""",
|
|
MATH_LIBS + SERVER_LIBS,
|
|
)
|
|
self.assertTrue(loaded["sympy"] and loaded["scipy"], "workers should preload the math libraries")
|
|
self.assertFalse(any(loaded[m] for m in ("mcp", "uvicorn", "starlette")), loaded)
|
|
self.assertTrue(loaded["single_threaded_blas"])
|
|
|
|
|
|
class MemoryLimitTest(unittest.TestCase):
|
|
def test_limit_is_relative_to_loaded_size(self):
|
|
"""RLIMIT_AS = current address space + configured extra (Linux path, simulated)."""
|
|
with (
|
|
mock.patch.object(builtins, "open", lambda *a, **k: io.StringIO("100000 5000 0 0 0 0 0")),
|
|
mock.patch.object(executor_module.resource, "setrlimit") as setrlimit,
|
|
):
|
|
executor_module._limit_address_space(512)
|
|
expected = 100000 * os.sysconf("SC_PAGE_SIZE") + 512 * 1024 * 1024
|
|
setrlimit.assert_called_once_with(executor_module.resource.RLIMIT_AS, (expected, expected))
|
|
|
|
def test_failure_is_not_fatal(self):
|
|
with mock.patch.object(builtins, "open", side_effect=OSError("no /proc")), self.assertLogs("mcp_calculator.executor", "WARNING") as logs:
|
|
executor_module._limit_address_space(512)
|
|
self.assertIn("could not set memory limit", logs.output[0])
|
|
|
|
|
|
class TaskDispatchTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_only_public_ops_functions_can_be_called(self):
|
|
executor = Executor(load_settings([], env={"MCP_CALC_EXECUTOR": "inline"}))
|
|
for task in ("os.system", "algebra._equivalent", "algebra", "parsing.parse", "algebra.__class__"):
|
|
with self.subTest(task=task), self.assertRaisesRegex(ComputationError, "unknown task"):
|
|
await executor.run(task)
|
|
|
|
|
|
class SpecConsistencyTest(unittest.TestCase):
|
|
def test_distribution_tables_are_consistent(self):
|
|
names = set(typing.get_args(spec.Distribution))
|
|
self.assertEqual(names, set(spec.DISTRIBUTION_PARAMETERS))
|
|
self.assertEqual(names, set(stats._DISTRIBUTIONS))
|