business-assistant/src/step4_api/tests/test_memory_cgroup_v1.py

"""Synthetic proc/cgroup fixtures; no host configuration or external calls."""

from mmap import PAGESIZE
from pathlib import Path
from tempfile import TemporaryDirectory
from types import SimpleNamespace
import sys
import unittest
from unittest.mock import Mock, patch

import httpx

from src.step4_api.app import create_app
from src.step4_api.config import Settings
from src.step4_api.memory import Admission, MIB, cgroup_memory, read_memory, require_hard_limit


class V1Fixture:
    def __init__(self, directory, group='/user.slice/user-1000.slice/session-1.scope', root='/'):
        self.proc = Path(directory) / 'proc'
        (self.proc / 'self').mkdir(parents=True)
        self.group_file = self.proc / 'self/cgroup'
        self.mountinfo = self.proc / 'self/mountinfo'
        self.mount = Path(directory) / 'memory controller'
        self.mount.mkdir()
        self.root = root
        self.group_file.write_text(f'12:rdma:/\n6:memory:{group}\n1:name=systemd:/session\n')
        self.write_mountinfo()
        relative = group[len(root):] if root != '/' and group.startswith(root) else group
        self.current = self.mount / relative.lstrip('/')
        self.current.mkdir(parents=True, exist_ok=True)
        node = self.current
        while True:
            self.set_memory(node, sys.maxsize // PAGESIZE * PAGESIZE, 0)
            if node == self.mount:
                break
            node = node.parent
        (self.proc / 'meminfo').write_text('MemTotal: 16777216 kB\nMemAvailable: 8388608 kB\n')
        (self.proc / 'self/statm').write_text(f'999999 {128 * MIB // PAGESIZE} 0\n')

    def write_mountinfo(self):
        mount = str(self.mount).replace('\\', '/').replace(' ', r'\040')
        self.mountinfo.write_text(f'33 22 0:29 {self.root} {mount} rw,nosuid - cgroup cgroup rw,memory\n')

    @staticmethod
    def set_memory(node, maximum, usage):
        (node / 'memory.limit_in_bytes').write_text(str(maximum))
        (node / 'memory.usage_in_bytes').write_text(str(usage))

    def read_cgroup(self, **kwargs):
        return cgroup_memory(self.group_file, mountinfo=self.mountinfo, **kwargs)


class CgroupV1Tests(unittest.TestCase):
    def fixture(self, **kwargs):
        temp = TemporaryDirectory()
        self.addCleanup(temp.cleanup)
        return V1Fixture(temp.name, **kwargs)

    def test_standalone_diagnostic_observes_parent_pressure(self):
        tree = self.fixture()
        tree.set_memory(tree.current, 2048 * MIB, 256 * MIB)
        tree.set_memory(tree.current.parent, 4096 * MIB, 4000 * MIB)
        self.assertEqual(tree.read_cgroup(), (2048 * MIB, 96 * MIB))
        with patch('src.step4_api.memory.sys.platform', 'linux'):
            self.assertEqual(read_memory(tree.proc).available, 96 * MIB)
            tree.set_memory(tree.current.parent, 4096 * MIB, 512 * MIB)
            self.assertEqual(read_memory(tree.proc).available, 1792 * MIB)

    def test_unlimited_uses_host_available_and_has_no_finite_hard_limit(self):
        tree = self.fixture()
        self.assertEqual(tree.read_cgroup(), (None, None))
        with patch('src.step4_api.memory.sys.platform', 'linux'):
            sample = read_memory(tree.proc)
        self.assertEqual((sample.rss, sample.available), (128 * MIB, 8192 * MIB))

    def test_root_group_and_over_limit_usage(self):
        tree = self.fixture(group='/')
        tree.set_memory(tree.mount, 100 * MIB, 101 * MIB)
        self.assertEqual(tree.read_cgroup(), (100 * MIB, 0))

    def test_host_headroom_still_limits_large_cgroup(self):
        tree = self.fixture()
        tree.set_memory(tree.current, 16 * 1024 * MIB, MIB)
        with patch('src.step4_api.memory.sys.platform', 'linux'):
            self.assertEqual(read_memory(tree.proc).available, 8192 * MIB)

    def test_combined_controller_and_hybrid_prefers_memory_v1(self):
        tree = self.fixture(group='/service')
        tree.group_file.write_text('0::/unified\n6:cpu,memory:/service\n')
        tree.set_memory(tree.current, 900 * MIB, 100 * MIB)
        self.assertEqual(tree.read_cgroup(), (900 * MIB, 800 * MIB))

    def test_bind_mount_root_and_namespaced_paths(self):
        for group in ('/tenant/service', '/service', '/'):
            with self.subTest(group=group):
                tree = self.fixture(group=group, root='/tenant')
                tree.set_memory(tree.current, 1500 * MIB, 200 * MIB)
                self.assertEqual(tree.read_cgroup(), (1500 * MIB, 1300 * MIB))

    def test_unrelated_mounts_are_ignored(self):
        tree = self.fixture()
        prefix = '44 22 0:30 / /missing rw - cgroup cgroup rw,cpu\n'
        tree.mountinfo.write_text(prefix + tree.mountinfo.read_text())
        self.assertEqual(tree.read_cgroup(), (None, None))

    def test_missing_mount_or_group_is_not_host_only_fallback(self):
        for kind in ('mount', 'group', 'controller', 'traversal'):
            with self.subTest(kind=kind):
                tree = self.fixture()
                if kind == 'mount':
                    tree.mountinfo.write_text('')
                elif kind == 'group':
                    tree.group_file.write_text('6:memory:/missing\n')
                elif kind == 'controller':
                    tree.group_file.write_text('6:cpu:/\n')
                else:
                    tree.group_file.write_text('6:memory:/../outside\n')
                with self.assertRaises(OSError):
                    tree.read_cgroup()

    def test_standalone_diagnostic_reports_corrupt_missing_or_denied_accounting(self):
        for kind in ('missing', 'bad_limit', 'negative_limit', 'negative_usage', 'denied'):
            with self.subTest(kind=kind):
                tree = self.fixture()
                path = tree.current / 'memory.usage_in_bytes'
                if kind == 'missing':
                    path.unlink()
                elif kind == 'bad_limit':
                    (tree.current / 'memory.limit_in_bytes').write_text('invalid')
                elif kind == 'negative_limit':
                    tree.set_memory(tree.current, -1, 0)
                elif kind == 'negative_usage':
                    tree.set_memory(tree.current, 100, -1)
                with patch('src.step4_api.memory.sys.platform', 'linux'):
                    if kind == 'denied':
                        with patch.object(Path, 'read_text', side_effect=PermissionError()):
                            with self.assertRaises(PermissionError):
                                read_memory(tree.proc)
                    else:
                        with self.assertRaises((OSError, ValueError)):
                            read_memory(tree.proc)

    def test_v1_sampling_does_not_claim_v2_service_isolation(self):
        tree = self.fixture()
        tree.set_memory(tree.current, 2048 * MIB, 100 * MIB)
        self.assertEqual(tree.read_cgroup(), (2048 * MIB, 1948 * MIB))
        with patch.dict('os.environ', {'API_REQUIRE_MEMORY_CGROUP': '1', 'API_MEMORY_HARD_LIMIT_MB': '4096'}), \
                patch('src.step4_api.memory.sys.platform', 'linux'), \
                patch('src.step4_api.memory.cgroup_memory', side_effect=tree.read_cgroup):
            with self.assertRaisesRegex(OSError, 'cgroup_v2_required'):
                require_hard_limit()


class V1HttpTests(unittest.IsolatedAsyncioTestCase):
    async def test_questions_run_without_memory_sampling_even_when_files_break(self):
        with TemporaryDirectory() as directory:
            tree = V1Fixture(directory)
            reader = Mock(side_effect=lambda: read_memory(tree.proc))
            admission = Admission(Settings(), reader)

            def run(thread_id, question, emit, *args, **kwargs):
                emit('result', {'response': 'synthetic answer'})
                emit('done', {})

            runtime = SimpleNamespace(admission=admission, close=lambda: None, run=Mock(side_effect=run))
            app = create_app(Settings(), lambda _: runtime)
            with patch('src.step4_api.memory.sys.platform', 'linux'):
                async with app.router.lifespan_context(app):
                    async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url='http://test') as client:
                        questions = ('我想在青浦开一家餐饮店，需要办理哪些手续？', '个人租房有什么支持？', '你好')
                        for i, question in enumerate(questions):
                            response = await client.post('/api/chat', json={'thread_id': f'v1-{i}', 'question': question})
                            self.assertEqual(response.status_code, 200)
                            self.assertIn('event: result', response.text)
                        self.assertEqual(runtime.run.call_count, len(questions))
                        tree.set_memory(tree.current, 400 * MIB, 350 * MIB)
                        response = await client.post('/api/chat', json={'thread_id': 'low', 'question': 'test'})
                        self.assertEqual(response.status_code, 200)
                        self.assertIn('event: result', response.text)
                        (tree.current / 'memory.limit_in_bytes').unlink()
                        response = await client.post('/api/chat', json={'thread_id': 'failed', 'question': 'test'})
                        self.assertEqual(response.status_code, 200)
                        self.assertIn('event: result', response.text)
                        self.assertEqual(runtime.run.call_count, len(questions) + 2)
                        reader.assert_not_called()