jundot--omlx/tests/test_process_memory_enforcer.py

# SPDX-License-Identifier: Apache-2.0
"""Tests for ProcessMemoryEnforcer."""

import asyncio
from contextlib import suppress
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

import omlx.process_memory_enforcer as pme
import omlx.utils.psutil_compat as psutil_compat
from omlx.process_memory_enforcer import ProcessMemoryEnforcer


def _make_enforcer(
    engine_pool,
    ceiling: int = 10 * 1024**3,
    tier: str = "balanced",
    poll_interval: float = 0.1,
    soft_threshold: float = 1.0,
    hard_threshold: float = 1.0,
    breakdown: dict | None = None,
    **kwargs,
) -> ProcessMemoryEnforcer:
    """Build an enforcer with a deterministic hard ceiling.

    The new enforcer derives its ceiling from system_memory + tier +
    live available memory, which is impractical to mock per
    test. We replace `_get_hard_limit_bytes` AND `_get_ceiling_breakdown`
    with constants so tests can exercise the watermark logic without
    juggling system mocks. Pass ``breakdown`` to distinguish the three
    component ceilings (static/dynamic/metal_cap) for propagation tests;
    omit to use the same ``ceiling`` for all three (equivalent old
    behavior).

    Passing `ceiling=0` disables the limit.
    """
    enforcer = ProcessMemoryEnforcer(
        engine_pool=engine_pool,
        memory_guard_tier=tier,
        poll_interval=poll_interval,
        soft_threshold=soft_threshold,
        hard_threshold=hard_threshold,
        **kwargs,
    )
    enforcer._soft_threshold = soft_threshold
    enforcer._get_hard_limit_bytes = lambda: int(ceiling)
    if breakdown is None:
        breakdown = {
            "static": int(ceiling),
            "dynamic": int(ceiling),
            "metal_cap": int(ceiling),
            "hard_limit": int(ceiling),
        }
    enforcer._get_ceiling_breakdown = lambda: dict(breakdown)
    return enforcer


def _cycling(values):
    """side_effect helper: yield each value, then repeat the last forever.

    Lets tests express the meaningful sequence of mocked memory values
    without having to count exact call sites in _check_and_enforce (the
    new 2-watermark path re-reads phys_footprint after eviction).
    """
    if not values:
        raise ValueError("need at least one value")
    state = {"i": 0}

    def _next(*_args, **_kwargs):
        i = state["i"]
        if i < len(values) - 1:
            state["i"] = i + 1
        return values[i]

    return _next


def _make_entry(model_id, engine=None, is_loading=False, is_pinned=False):
    """Create a mock EngineEntry."""
    entry = MagicMock()
    entry.model_id = model_id
    entry.engine = engine
    entry.is_loading = is_loading
    entry.is_pinned = is_pinned
    entry.abort_loading = False
    entry.in_use = 0
    entry.last_access = 0.0
    entry.pending_unload_reason = None
    entry.abort_requested = False
    return entry


def _close_coro(coro):
    """side_effect for patched asyncio.create_task that closes the coroutine.

    Tests that mock create_task to skip the background loop still pass the
    real `_enforcement_loop()` coroutine in. Closing it prevents the
    "coroutine was never awaited" RuntimeWarning at gc time.
    """
    if hasattr(coro, "close"):
        coro.close()
    return MagicMock()


class TestMacOSVMStats:
    """Tests for the host_statistics64 telemetry adapter."""

    def test_uses_max_sized_host_info64_buffer(self):
        """Newer macOS kernels can require a larger vm_statistics64 tail."""

        class FakeLibc:
            def host_statistics64(self, host, flavor, stats, count):
                assert host == 123
                assert flavor == psutil_compat._HOST_VM_INFO64
                assert count._obj.value == psutil_compat._HOST_INFO64_MAX_COUNT
                stats[0] = 10
                stats[1] = 20
                stats[2] = 30
                stats[3] = 40
                count._obj.value = 104
                return 0

        with (
            patch.object(psutil_compat, "_libc", FakeLibc()),
            patch.object(psutil_compat, "_MACH_HOST", 123),
            patch.object(psutil_compat, "_VM_PAGE_SIZE", 4096),
        ):
            stats = pme.get_macos_vm_stats()

        assert stats == {
            "free": 10 * 4096,
            "active": 20 * 4096,
            "inactive": 30 * 4096,
            "wired": 40 * 4096,
        }

    def test_short_host_info64_response_returns_none(self):
        class FakeLibc:
            def host_statistics64(self, host, flavor, stats, count):
                count._obj.value = 3
                return 0

        with (
            patch.object(psutil_compat, "_libc", FakeLibc()),
            patch.object(psutil_compat, "_MACH_HOST", 123),
        ):
            assert pme.get_macos_vm_stats() is None


@pytest.fixture
def mock_engine_pool():
    """Create a mock EnginePool with required methods."""
    pool = MagicMock()
    pool._lock = asyncio.Lock()
    pool._find_lru_victim = MagicMock(return_value="model-a")
    pool._unload_engine = AsyncMock()
    pool._entries = {}

    def _entry_busy(entry):
        if getattr(entry, "in_use", 0) > 0:
            return True
        engine = getattr(entry, "engine", None)
        has_active = getattr(engine, "has_active_requests", None)
        if callable(has_active):
            return has_active() is True
        return False

    def _find_pending_unload_ready_locked():
        candidates = []
        for mid, entry in pool._entries.items():
            if not getattr(entry, "pending_unload_reason", None):
                continue
            if (
                getattr(entry, "engine", None) is None
                or getattr(entry, "is_loading", False)
                or getattr(entry, "is_pinned", False)
                or _entry_busy(entry)
            ):
                continue
            candidates.append((getattr(entry, "last_access", 0.0), mid))
        if not candidates:
            return None
        candidates.sort()
        return candidates[0][1]

    async def _unload_pending_if_idle_locked(model_id):
        entry = pool._entries.get(model_id)
        if (
            entry is None
            or getattr(entry, "engine", None) is None
            or not getattr(entry, "pending_unload_reason", None)
            or getattr(entry, "is_loading", False)
            or getattr(entry, "is_pinned", False)
            or _entry_busy(entry)
        ):
            return False
        entry.pending_unload_reason = None
        entry.abort_requested = False
        await pool._unload_engine(model_id)
        return True

    def _mark_pending_unload_locked(model_id, reason, *, abort_requested=False):
        entry = pool._entries.get(model_id)
        if (
            entry is None
            or getattr(entry, "engine", None) is None
            or getattr(entry, "is_loading", False)
            or getattr(entry, "is_pinned", False)
        ):
            return False
        entry.pending_unload_reason = reason
        if abort_requested:
            entry.abort_requested = True
        return True

    pool._find_pending_unload_ready_locked = MagicMock(
        side_effect=_find_pending_unload_ready_locked
    )
    pool._unload_pending_if_idle_locked = AsyncMock(
        side_effect=_unload_pending_if_idle_locked
    )
    pool._mark_pending_unload_locked = MagicMock(
        side_effect=_mark_pending_unload_locked
    )
    return pool


@pytest.fixture
def enforcer(mock_engine_pool):
    """Create an enforcer with a fixed 10GB ceiling.

    Soft/hard thresholds default to 1.0 so legacy single-threshold tests
    keep treating the ceiling as the single trip point. Dedicated
    2-watermark tests construct their own enforcer with default thresholds.
    """
    return _make_enforcer(mock_engine_pool, ceiling=10 * 1024**3)


class TestAdaptivePolling:
    """Tests for adaptive ProcessMemoryEnforcer polling cadence."""

    def test_no_loaded_or_loading_models_uses_unloaded_idle_interval(self, enforcer):
        enforcer._engine_pool._entries = {
            f"model-{i}": _make_entry(f"model-{i}") for i in range(3)
        }

        assert enforcer._select_poll_interval() == 30.0

    def test_loading_model_uses_active_interval(self, enforcer):
        enforcer._engine_pool._entries = {
            "loading": _make_entry("loading", is_loading=True)
        }

        assert enforcer._select_poll_interval() == enforcer._active_poll_interval

    def test_loaded_idle_model_uses_loaded_idle_interval(self, enforcer):
        engine = MagicMock()
        engine.has_active_requests.return_value = False
        enforcer._engine_pool._entries = {
            "loaded": _make_entry("loaded", engine=engine)
        }

        assert enforcer._select_poll_interval() == 10.0

    def test_loaded_active_model_uses_active_interval(self, enforcer):
        engine = MagicMock()
        engine.has_active_requests.return_value = True
        enforcer._engine_pool._entries = {
            "active": _make_entry("active", engine=engine)
        }

        assert enforcer._select_poll_interval() == enforcer._active_poll_interval

    def test_pressure_uses_active_interval(self, enforcer):
        enforcer._pressure_level = "soft"

        assert enforcer._select_poll_interval() == enforcer._active_poll_interval

    def test_activity_hint_uses_active_interval(self, enforcer):
        enforcer._engine_pool._entries = {}
        enforcer.wake(active=True)

        assert enforcer._select_poll_interval() == enforcer._active_poll_interval

    @pytest.mark.asyncio
    async def test_wake_interrupts_idle_sleep(self, enforcer):
        enforcer._engine_pool._entries = {}
        enforcer._running = True
        enforcer._check_and_enforce = AsyncMock()
        enforcer._check_ttl = AsyncMock()

        task = asyncio.create_task(enforcer._enforcement_loop())
        try:
            for _ in range(20):
                if enforcer._check_and_enforce.await_count >= 1:
                    break
                await asyncio.sleep(0.01)
            assert enforcer._check_and_enforce.await_count >= 1

            enforcer.wake()
            for _ in range(20):
                if enforcer._check_and_enforce.await_count >= 2:
                    break
                await asyncio.sleep(0.01)

            assert enforcer._check_and_enforce.await_count >= 2
        finally:
            enforcer._running = False
            enforcer.wake()
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task


class TestCheckAndEnforce:
    """Tests for _check_and_enforce method."""

    @pytest.mark.asyncio
    async def test_no_action_when_under_limit(self, enforcer):
        """No eviction when memory is under limit."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 5 * 1024**3
            await enforcer._check_and_enforce()
        enforcer._engine_pool._unload_engine.assert_not_called()

    @pytest.mark.asyncio
    async def test_no_action_at_exact_limit(self, enforcer):
        """No eviction when memory is exactly at limit."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 10 * 1024**3
            await enforcer._check_and_enforce()
        enforcer._engine_pool._unload_engine.assert_not_called()

    @pytest.mark.asyncio
    async def test_evicts_when_over_limit(self, enforcer):
        """Evicts LRU model when over limit (multiple models loaded)."""
        # Need at least 2 loaded non-pinned models for eviction path
        engine_a = MagicMock()
        engine_a.abort_all_requests = AsyncMock(return_value=0)
        engine_b = MagicMock()
        engine_b.abort_all_requests = AsyncMock(return_value=0)
        entry_a = _make_entry("model-a", engine=engine_a)
        entry_b = _make_entry("model-b", engine=engine_b)
        enforcer._engine_pool._entries = {
            "model-a": entry_a,
            "model-b": entry_b,
        }
        enforcer._engine_pool._find_lru_victim.return_value = "model-a"

        async def fake_unload(model_id):
            enforcer._engine_pool._entries[model_id].engine = None

        enforcer._engine_pool._unload_engine.side_effect = fake_unload

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    15 * 1024**3,  # Initial check (over limit)
                    15 * 1024**3,  # Re-check before eviction loop
                    8 * 1024**3,  # After eviction (under limit)
                ]
            )
            await enforcer._check_and_enforce()
        enforcer._engine_pool._unload_engine.assert_called_once_with("model-a")

    @pytest.mark.asyncio
    async def test_stops_when_all_pinned(self, enforcer):
        """Stops eviction when all models are pinned (no victim)."""
        enforcer._engine_pool._find_lru_victim.return_value = None
        # Add a pinned loaded model so the log says "pinned"
        engine = MagicMock()
        engine.abort_all_requests = AsyncMock(return_value=3)
        entry = _make_entry("pinned-model", engine=engine, is_pinned=True)
        enforcer._engine_pool._entries = {"pinned-model": entry}
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    11 * 1024**3,  # Initial check, over ceiling but not emergency
                    11 * 1024**3,  # Re-check in loop
                ]
            )
            await enforcer._check_and_enforce()
        enforcer._engine_pool._unload_engine.assert_not_called()
        engine.abort_all_requests.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_evicts_multiple_models(self, enforcer):
        """Evicts multiple models in sequence until under limit."""
        # Need 3 loaded non-pinned models for sequential eviction
        engine_a = MagicMock()
        engine_a.abort_all_requests = AsyncMock(return_value=0)
        engine_b = MagicMock()
        engine_b.abort_all_requests = AsyncMock(return_value=0)
        engine_c = MagicMock()
        engine_c.abort_all_requests = AsyncMock(return_value=0)
        entry_a = _make_entry("model-a", engine=engine_a)
        entry_b = _make_entry("model-b", engine=engine_b)
        entry_c = _make_entry("model-c", engine=engine_c)
        enforcer._engine_pool._entries = {
            "model-a": entry_a,
            "model-b": entry_b,
            "model-c": entry_c,
        }
        enforcer._engine_pool._find_lru_victim.side_effect = [
            "model-a",
            "model-b",
        ]

        async def fake_unload(model_id):
            enforcer._engine_pool._entries[model_id].engine = None

        enforcer._engine_pool._unload_engine.side_effect = fake_unload

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    20 * 1024**3,  # Initial check
                    20 * 1024**3,  # Re-check (still over)
                    15 * 1024**3,  # After first eviction (still over)
                    8 * 1024**3,  # After second eviction (under limit)
                ]
            )
            await enforcer._check_and_enforce()
        assert enforcer._engine_pool._unload_engine.call_count == 2

    @pytest.mark.asyncio
    async def test_aborts_loading_model_when_no_lru_victim(self, enforcer):
        """Aborts a loading model when no LRU victim is available."""
        enforcer._engine_pool._find_lru_victim.return_value = None
        loading_entry = _make_entry("loading-model", engine=None, is_loading=True)
        enforcer._engine_pool._entries = {"loading-model": loading_entry}

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    15 * 1024**3,  # Initial check
                    15 * 1024**3,  # Re-check in loop
                ]
            )
            await enforcer._check_and_enforce()

        assert loading_entry.abort_loading is True
        enforcer._engine_pool._unload_engine.assert_not_called()

    @pytest.mark.asyncio
    async def test_evicts_lru_before_aborting_loading(self, enforcer):
        """Evicts LRU models first, then aborts loading model."""
        # Need 2 loaded non-pinned so model-a gets evicted (not abort path)
        engine_a = MagicMock()
        engine_a.abort_all_requests = AsyncMock(return_value=0)
        engine_b = MagicMock()
        engine_b.abort_all_requests = AsyncMock(return_value=0)
        entry_a = _make_entry("model-a", engine=engine_a)
        entry_b = _make_entry("model-b", engine=engine_b)
        loading_entry = _make_entry("loading-model", engine=None, is_loading=True)
        enforcer._engine_pool._entries = {
            "model-a": entry_a,
            "model-b": entry_b,
            "loading-model": loading_entry,
        }

        async def fake_unload(model_id):
            enforcer._engine_pool._entries[model_id].engine = None

        enforcer._engine_pool._unload_engine.side_effect = fake_unload

        # First call returns victim, second call returns None
        enforcer._engine_pool._find_lru_victim.side_effect = [
            "model-a",
            None,
        ]

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    20 * 1024**3,  # Initial check
                    20 * 1024**3,  # Re-check (still over)
                    15 * 1024**3,  # After eviction (still over)
                ]
            )
            await enforcer._check_and_enforce()

        # LRU victim evicted first
        enforcer._engine_pool._unload_engine.assert_called_once_with("model-a")
        # Then loading model abort requested
        assert loading_entry.abort_loading is True

    @pytest.mark.asyncio
    async def test_no_models_loaded_or_loading(self, enforcer):
        """Logs correctly when no models are loaded or loading."""
        enforcer._engine_pool._find_lru_victim.return_value = None
        enforcer._engine_pool._entries = {}

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    15 * 1024**3,  # Initial check
                    15 * 1024**3,  # Re-check
                ]
            )
            await enforcer._check_and_enforce()
        # Should not raise, just log warning


class TestCurrentUsageTelemetry:
    """Tests for enforcer memory telemetry threading behavior."""

    def test_idle_usage_keeps_direct_mlx_telemetry(self, enforcer):
        """Idle accounting preserves the legacy max(active, phys_footprint)."""
        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=3 * 1024**3,
            ),
        ):
            mock_mx.get_active_memory.return_value = 5 * 1024**3
            assert enforcer._current_usage_bytes() == 5 * 1024**3
        mock_mx.get_active_memory.assert_called_once()

    def test_active_usage_uses_cached_executor_sample_without_mlx(self, enforcer):
        """Active decode ticks must not call MLX from the enforcer thread."""
        scheduler = MagicMock()
        scheduler.get_cached_mlx_active_memory_bytes.return_value = 7 * 1024**3
        engine = MagicMock()
        engine.has_active_requests.return_value = True
        engine.scheduler = scheduler
        enforcer._engine_pool._entries = {"m": _make_entry("m", engine=engine)}

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=5 * 1024**3,
            ),
        ):
            mock_mx.get_active_memory.side_effect = AssertionError(
                "background enforcer touched MLX during active decode"
            )
            assert enforcer._current_usage_bytes() == 7 * 1024**3
        mock_mx.get_active_memory.assert_not_called()

    def test_active_usage_keeps_phys_footprint_when_it_dominates(self, enforcer):
        scheduler = MagicMock()
        scheduler.get_cached_mlx_active_memory_bytes.return_value = 4 * 1024**3
        engine = MagicMock()
        engine.has_active_requests.return_value = True
        engine.scheduler = scheduler
        enforcer._engine_pool._entries = {"m": _make_entry("m", engine=engine)}

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=9 * 1024**3,
            ),
        ):
            assert enforcer._current_usage_bytes() == 9 * 1024**3
        mock_mx.get_active_memory.assert_not_called()

    def test_hard_limit_reuses_cached_metal_cap(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="custom",
            memory_guard_custom_ceiling_gb=1024.0,
        )
        enforcer._effective_metal_cap_bytes = 48 * 1024**3
        with (
            patch("omlx.settings.get_system_memory", return_value=64 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                side_effect=AssertionError("Metal cap cache was bypassed"),
            ),
        ):
            assert enforcer._get_hard_limit_bytes() == 48 * 1024**3


class TestDisabledWhenCeilingZero:
    """Tests for enforcement disabled when the ceiling is 0 (guard off)."""

    @pytest.mark.asyncio
    async def test_no_enforce_when_ceiling_zero(self, mock_engine_pool):
        """No enforcement when ceiling is 0 (guard disabled)."""
        enforcer = _make_enforcer(mock_engine_pool, ceiling=0)
        engine = MagicMock()
        engine.abort_all_requests = AsyncMock(return_value=0)
        entry = _make_entry("model-a", engine=engine)
        mock_engine_pool._entries = {"model-a": entry}

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            await enforcer._check_and_enforce()

        engine.abort_all_requests.assert_not_awaited()
        mock_engine_pool._unload_engine.assert_not_called()

    @pytest.mark.asyncio
    async def test_no_enforce_when_guard_off(self, mock_engine_pool):
        """No enforcement when memory guard toggle is off."""
        enforcer = _make_enforcer(mock_engine_pool, ceiling=0)
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            await enforcer._check_and_enforce()

        mock_engine_pool._unload_engine.assert_not_called()

    @pytest.mark.asyncio
    async def test_propagate_zero_disables_inline_prefill_check(self, mock_engine_pool):
        """Propagating ceiling=0 sets scheduler limit to 0 (disabled)."""
        enforcer = _make_enforcer(
            mock_engine_pool,
            ceiling=0,
            breakdown={
                "static": 0,
                "dynamic": 0,
                "metal_cap": 0,
                "hard_limit": 0,
            },
        )
        bg = MagicMock(spec=[])
        bg._memory_limit_bytes = 999
        bg._memory_hard_limit_bytes = 999
        scheduler = MagicMock(spec=[])
        scheduler._memory_limit_bytes = 999
        scheduler._memory_hard_limit_bytes = 999
        scheduler.batch_generator = bg
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        mock_engine_pool._entries = {"model-a": entry}

        enforcer._propagate_memory_limit()

        assert scheduler._memory_limit_bytes == 0
        assert scheduler._memory_hard_limit_bytes == 0
        assert bg._memory_limit_bytes == 0
        assert bg._memory_hard_limit_bytes == 0

    @pytest.mark.asyncio
    async def test_propagate_ceiling_components_to_scheduler(self, mock_engine_pool):
        """All three component ceilings + the tier name must reach the
        scheduler so the binding-aware rejection message has the inputs
        it needs to identify which knob the operator should turn."""
        static_b = 64 * 1024**3
        dynamic_b = 16 * 1024**3
        metal_b = 48 * 1024**3
        enforcer = _make_enforcer(
            mock_engine_pool,
            tier="custom",
            ceiling=min(static_b, dynamic_b, metal_b),
            breakdown={
                "static": static_b,
                "dynamic": dynamic_b,
                "metal_cap": metal_b,
                "hard_limit": min(static_b, dynamic_b, metal_b),
            },
        )
        scheduler = MagicMock(spec=[])
        scheduler._memory_limit_bytes = 0
        scheduler._memory_hard_limit_bytes = 0
        scheduler._memory_static_ceiling_bytes = 0
        scheduler._memory_dynamic_ceiling_bytes = 0
        scheduler._memory_metal_cap_bytes = 0
        scheduler._memory_guard_tier = "balanced"
        scheduler._prefill_headroom_safety = 0.0
        scheduler.batch_generator = None
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        mock_engine_pool._entries = {"model-a": entry}

        enforcer._propagate_memory_limit()

        assert scheduler._memory_static_ceiling_bytes == static_b
        assert scheduler._memory_dynamic_ceiling_bytes == dynamic_b
        assert scheduler._memory_metal_cap_bytes == metal_b
        assert scheduler._memory_hot_cache_reserved_bytes == 0
        assert scheduler._memory_guard_tier == "custom", (
            "tier name must reach the scheduler so the advice ladder can "
            "distinguish dynamic-on-custom (raise custom_ceiling_bytes) "
            "from dynamic-on-reclaim-tier (close other apps)"
        )
        assert scheduler._prefill_headroom_safety == 0.90
        assert (
            scheduler._memory_hard_limit_bytes == dynamic_b
        ), "hard limit must be min of the three components"

    @pytest.mark.asyncio
    async def test_propagates_hot_cache_reservation_for_binding_messages(
        self, mock_engine_pool
    ):
        """The scheduler hard limit subtracts hot-cache reservation, so the
        rejection formatter needs the same reservation to identify which
        original component is binding."""
        static_b = 64 * 1024**3
        dynamic_b = 32 * 1024**3
        metal_b = 16 * 1024**3
        hot_reserved_b = 2 * 1024**3
        enforcer = _make_enforcer(
            mock_engine_pool,
            ceiling=metal_b,
            breakdown={
                "static": static_b,
                "dynamic": dynamic_b,
                "metal_cap": metal_b,
                "hard_limit": metal_b,
            },
        )
        enforcer._hot_cache_reserved_bytes = lambda: hot_reserved_b
        scheduler = MagicMock(spec=[])
        scheduler._memory_limit_bytes = 0
        scheduler._memory_hard_limit_bytes = 0
        scheduler._memory_static_ceiling_bytes = 0
        scheduler._memory_dynamic_ceiling_bytes = 0
        scheduler._memory_metal_cap_bytes = 0
        scheduler._memory_hot_cache_reserved_bytes = 0
        scheduler.batch_generator = None
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        mock_engine_pool._entries = {"model-a": entry}

        enforcer._propagate_memory_limit()

        assert scheduler._memory_hard_limit_bytes == metal_b - hot_reserved_b
        assert scheduler._memory_static_ceiling_bytes == static_b
        assert scheduler._memory_dynamic_ceiling_bytes == dynamic_b
        assert scheduler._memory_metal_cap_bytes == metal_b
        assert scheduler._memory_hot_cache_reserved_bytes == hot_reserved_b


class TestPrefillMemoryGuardToggle:
    """Tests for prefill_memory_guard setter and Metal limit management."""

    def test_enable_guard_is_noop_for_metal_limits(self, enforcer):
        """Enabling guard does NOT call Metal limits (no-op since #429)."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            enforcer._running = True

            enforcer.prefill_memory_guard = True
            assert enforcer.prefill_memory_guard is True
            mock_mx.set_memory_limit.assert_not_called()
            mock_mx.set_cache_limit.assert_not_called()

    def test_disable_guard_is_noop_for_metal_limits(self, enforcer):
        """Disabling guard does NOT call Metal limits (no-op since #429)."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            enforcer._running = True

            enforcer.prefill_memory_guard = True
            enforcer.prefill_memory_guard = False
            assert enforcer.prefill_memory_guard is False
            mock_mx.set_memory_limit.assert_not_called()
            mock_mx.set_cache_limit.assert_not_called()

    def test_disable_guard_noop_without_prior_limits(self, enforcer):
        """Disabling guard when no limits were set does not call mx."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            enforcer._running = True

            # Disable without enabling first
            enforcer.prefill_memory_guard = False
            mock_mx.set_memory_limit.assert_not_called()
            mock_mx.set_cache_limit.assert_not_called()


class TestStaticCeiling:
    """Tier-driven static ceiling (`total_ram - tier.static_reserve`).

    >= 24 GB systems use a tier-scaled reserve. < 24 GB systems always
    use a 4 GB reserve regardless of tier.
    """

    @pytest.mark.parametrize(
        "tier,expected_reserve_gb",
        [("safe", 8), ("balanced", 6), ("aggressive", 4)],
    )
    def test_large_system_tier_reserve(
        self, mock_engine_pool, tier, expected_reserve_gb
    ):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier=tier
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 96 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == (96 - expected_reserve_gb) * 1024**3

    @pytest.mark.parametrize("tier", ["safe", "balanced", "aggressive"])
    def test_small_system_uses_4gb_reserve_regardless_of_tier(
        self, mock_engine_pool, tier
    ):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier=tier
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 12 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == 8 * 1024**3

    @pytest.mark.parametrize("tier", ["safe", "balanced", "aggressive"])
    def test_16gb_system_uses_4gb_reserve_regardless_of_tier(
        self, mock_engine_pool, tier
    ):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier=tier
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 16 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == 12 * 1024**3

    @pytest.mark.parametrize("tier", ["safe", "balanced", "aggressive"])
    def test_between_16gb_and_24gb_system_uses_4gb_reserve_regardless_of_tier(
        self, mock_engine_pool, tier
    ):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier=tier
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 16 * 1024**3 + 256 * 1024**2
            result = enforcer._get_static_ceiling()
        assert result == 12 * 1024**3 + 256 * 1024**2

    def test_18gb_system_uses_4gb_reserve(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 18 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == 14 * 1024**3

    def test_24gb_system_uses_tier_reserve(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 24 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == 18 * 1024**3

    def test_custom_uses_2gb_reserve_on_large_system(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="custom",
            memory_guard_custom_ceiling_gb=50.0,
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 64 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == 62 * 1024**3

    def test_custom_uses_2gb_reserve_on_small_system(self, mock_engine_pool):
        """Custom bypasses the 4 GB small-system reserve."""
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="custom",
            memory_guard_custom_ceiling_gb=8.0,
        )
        with patch("omlx.settings.get_system_memory") as mock_mem:
            mock_mem.return_value = 12 * 1024**3
            result = enforcer._get_static_ceiling()
        assert result == 10 * 1024**3


class TestDynamicCeilingActiveRatio:
    """Dynamic ceiling sums free + inactive + active * tier ratio
    (host_statistics64 path) for safe / balanced / aggressive."""

    @pytest.mark.parametrize(
        "tier,ratio",
        [("safe", 0.2), ("balanced", 0.5), ("aggressive", 0.8)],
    )
    def test_active_ratio_per_tier(self, mock_engine_pool, tier, ratio):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier=tier
        )
        with (
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=1 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value={
                    "free": 10 * 1024**3,
                    "inactive": 4 * 1024**3,
                    "active": 8 * 1024**3,
                    "wired": 2 * 1024**3,
                },
            ),
        ):
            result = enforcer._get_dynamic_ceiling()
        expected = 1 * 1024**3 + 10 * 1024**3 + 4 * 1024**3 + int(8 * 1024**3 * ratio)
        assert result == expected

    def test_non_macos_falls_back_to_compat_available(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=2 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value=None,
            ),
            patch(
                "omlx.process_memory_enforcer.psutil_compat.virtual_memory",
                return_value=SimpleNamespace(available=15 * 1024**3),
            ) as mock_virtual_memory,
        ):
            result = enforcer._get_dynamic_ceiling()
        assert result == 2 * 1024**3 + 15 * 1024**3
        mock_virtual_memory.assert_called_once()

    def test_macos_vm_stat_failure_uses_compat_available(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=2 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value=None,
            ),
            patch(
                "omlx.process_memory_enforcer.psutil_compat.virtual_memory",
                return_value=SimpleNamespace(available=15 * 1024**3),
            ) as mock_virtual_memory,
        ):
            result = enforcer._get_dynamic_ceiling()

        assert result == 2 * 1024**3 + 15 * 1024**3
        mock_virtual_memory.assert_called_once()

    def test_compat_failure_falls_back_to_static_ceiling(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=2 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value=None,
            ),
            patch(
                "omlx.process_memory_enforcer.psutil_compat.virtual_memory",
                side_effect=RuntimeError(
                    "host_statistics64(HOST_VM_INFO64) syscall failed"
                ),
            ),
            patch("omlx.settings.get_system_memory", return_value=64 * 1024**3),
        ):
            result = enforcer._get_dynamic_ceiling()

        assert result == 58 * 1024**3


class TestDynamicCeilingCustom:
    """tier == custom: dynamic is the user-specified value, panic-safe via min()."""

    def test_custom_ceiling_returned_verbatim(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="custom",
            memory_guard_custom_ceiling_gb=24.0,
        )
        assert enforcer._get_dynamic_ceiling() == 24 * 1024**3

    def test_custom_setter_updates_ceiling(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="custom",
            memory_guard_custom_ceiling_gb=10.0,
        )
        enforcer.memory_guard_custom_ceiling_bytes = 30 * 1024**3
        assert enforcer._get_dynamic_ceiling() == 30 * 1024**3

    def test_custom_clamped_by_static_and_metal_cap(self, mock_engine_pool):
        """User can type any number; the final ceiling is still min(static,
        dynamic, metal_cap) so out-of-range values are panic-safe."""
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="custom",
            memory_guard_custom_ceiling_gb=1024.0,  # absurdly large
        )
        with (
            patch("omlx.settings.get_system_memory", return_value=64 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=48 * 1024**3,
            ),
        ):
            ceiling = enforcer._get_hard_limit_bytes()
        # static = 64 - 2 = 62 GB; metal = 48 GB; custom = 1024 GB
        # → clamped to metal 48 GB
        assert ceiling == 48 * 1024**3


class TestHardLimitCalculation:
    """`_get_hard_limit_bytes` returns min(static, dynamic), or 0 when guard off."""

    def test_picks_static_when_smaller(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch("omlx.settings.get_system_memory") as mock_mem,
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=2 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value={
                    "free": 30 * 1024**3,
                    "inactive": 10 * 1024**3,
                    "active": 5 * 1024**3,
                    "wired": 1 * 1024**3,
                },
            ),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=100 * 1024**3,
            ),
        ):
            mock_mem.return_value = 48 * 1024**3  # static = 42 GB
            # dynamic balanced = 2 + 30 + 10 + 5*0.5 = 44.5 GB
            # static (42) wins → final ceiling is 42 GB
            assert enforcer._get_hard_limit_bytes() == 42 * 1024**3

    def test_picks_dynamic_when_smaller(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch("omlx.settings.get_system_memory") as mock_mem,
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=1 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value={
                    "free": 5 * 1024**3,
                    "inactive": 2 * 1024**3,
                    "active": 4 * 1024**3,
                    "wired": 1 * 1024**3,
                },
            ),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=100 * 1024**3,
            ),
        ):
            mock_mem.return_value = 48 * 1024**3  # static = 42 GB
            # dynamic balanced = 1 + 5 + 2 + int(4 * 0.5) = 10 GB
            # → dynamic wins
            assert enforcer._get_hard_limit_bytes() == 10 * 1024**3


class TestAbortLimitCalculation:
    """`_get_abort_limit_bytes` = min(static, metal_cap); ignores the jittery
    dynamic ceiling so a transient dip can't kill an in-flight prefill."""

    def test_ignores_dynamic_ceiling(self, mock_engine_pool):
        """Even when the dynamic ceiling is tiny, the abort limit stays at the
        stable min(static, metal_cap)."""
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch("omlx.settings.get_system_memory") as mock_mem,
            patch(
                "omlx.process_memory_enforcer.get_phys_footprint",
                return_value=1 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_macos_vm_stats",
                return_value={  # dynamic would compute to ~10 GB (depressed)
                    "free": 5 * 1024**3,
                    "inactive": 2 * 1024**3,
                    "active": 4 * 1024**3,
                    "wired": 1 * 1024**3,
                },
            ),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=100 * 1024**3,
            ),
        ):
            mock_mem.return_value = 48 * 1024**3  # static = 42 GB
            # dynamic (~10 GB) is far below, but the abort limit ignores it.
            assert enforcer._get_abort_limit_bytes() == 42 * 1024**3
            # Sanity: the (jittery) hard limit DID drop to dynamic.
            assert enforcer._get_hard_limit_bytes() == 10 * 1024**3

    def test_picks_metal_cap_when_smaller_than_static(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch("omlx.settings.get_system_memory", return_value=64 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=43 * 1024**3,
            ),
        ):
            # static = 64 - 8 = 56 GB; metal = 43 GB → min = 43 GB
            assert enforcer._get_abort_limit_bytes() == 43 * 1024**3

    def test_falls_back_to_static_when_metal_cap_unknown(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="aggressive"
        )
        with (
            patch("omlx.settings.get_system_memory", return_value=48 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=0,  # unknown
            ),
        ):
            # aggressive reserve = 4 GB → static = 44 GB
            assert enforcer._get_abort_limit_bytes() == 44 * 1024**3

    def test_abort_limit_zero_when_guard_disabled(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, prefill_memory_guard=False
        )
        assert enforcer._get_abort_limit_bytes() == 0

    def test_hard_limit_zero_when_guard_disabled(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="balanced",
            prefill_memory_guard=False,
        )
        assert enforcer._get_hard_limit_bytes() == 0
        assert enforcer.get_final_ceiling() == 0

    def test_unknown_tier_falls_back_to_balanced(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="extreme"
        )
        assert enforcer.memory_guard_tier == "balanced"


class TestMetalWiredLimit:
    """enforcer.start() applies MLX wired limits only for explicit sysctl caps."""

    def test_start_calls_set_wired_limit_with_static_ceiling(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch("omlx.settings.get_system_memory", return_value=48 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=64 * 1024**3,  # cap above static ceiling
            ),
            patch(
                "omlx.process_memory_enforcer.get_iogpu_wired_limit_bytes",
                return_value=64 * 1024**3,
            ),
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch.object(asyncio, "create_task", side_effect=_close_coro),
        ):
            mock_mx.set_wired_limit.return_value = 36 * 1024**3
            enforcer.start()
        # balanced @ 48 GB => static_ceiling = 42 GB
        mock_mx.set_wired_limit.assert_called_once_with(42 * 1024**3)
        # Stored value is the desired ceiling (not the post-clamp value)
        # so the admin UI can detect a kernel cap that's below the
        # request and surface the sysctl-raise hint.
        assert enforcer._metal_wired_limit_request == 42 * 1024**3

    def test_start_clamps_to_effective_cap_when_lower(self, mock_engine_pool):
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="aggressive"
        )
        with (
            patch("omlx.settings.get_system_memory", return_value=64 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=42 * 1024**3,  # cap < static ceiling
            ),
            patch(
                "omlx.process_memory_enforcer.get_iogpu_wired_limit_bytes",
                return_value=42 * 1024**3,
            ),
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch.object(asyncio, "create_task", side_effect=_close_coro),
        ):
            mock_mx.set_wired_limit.return_value = 48 * 1024**3
            enforcer.start()
        # aggressive @ 64 GB static = 60 GB, clamped to cap 42 GB
        mock_mx.set_wired_limit.assert_called_once_with(42 * 1024**3)
        # Desired (60 GB) is stored, not the post-clamp 42 GB.
        assert enforcer._metal_wired_limit_request == 60 * 1024**3

    def test_start_skips_set_wired_limit_when_sysctl_unset(
        self, mock_engine_pool, caplog
    ):
        """sysctl=0 path: plan against Apple cap, but don't touch MLX state."""
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            with (
                patch("omlx.settings.get_system_memory", return_value=512 * 1024**3),
                patch(
                    "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                    return_value=128 * 1024**3,  # Apple default below static ceiling
                ),
                patch(
                    "omlx.process_memory_enforcer.get_iogpu_wired_limit_bytes",
                    return_value=0,  # sysctl unset; cap comes from working set
                ),
                patch("omlx.process_memory_enforcer.mx") as mock_mx,
                patch.object(asyncio, "create_task", side_effect=_close_coro),
            ):
                enforcer.start()
        # balanced @ 512 GB static = 506 GB, clamped to the 5%-of-RAM
        # recommendation reserve for the admin banner (#2184). The scheduler
        # still clamps to the 128 GB effective cap, but MLX's wired limit is
        # left untouched.
        mock_mx.set_wired_limit.assert_not_called()
        total = 512 * 1024**3
        assert enforcer._metal_wired_limit_request == total - total // 20
        assert "leaving Apple's default Metal cap active" in caplog.text

    def test_start_handles_set_wired_limit_error(self, mock_engine_pool):
        """Older macOS (<15) raises on the call; enforcer keeps going."""
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool, memory_guard_tier="balanced"
        )
        with (
            patch("omlx.settings.get_system_memory", return_value=48 * 1024**3),
            patch(
                "omlx.process_memory_enforcer.get_effective_metal_cap_bytes",
                return_value=64 * 1024**3,
            ),
            patch(
                "omlx.process_memory_enforcer.get_iogpu_wired_limit_bytes",
                return_value=64 * 1024**3,
            ),
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch.object(asyncio, "create_task", side_effect=_close_coro),
        ):
            mock_mx.set_wired_limit.side_effect = RuntimeError("unsupported")
            enforcer.start()  # must not raise
        # We store the desired static_ceiling even when the call fails,
        # so the admin UI can still surface a warning.
        assert enforcer._metal_wired_limit_request == 42 * 1024**3

    def test_start_skips_when_guard_disabled(self, mock_engine_pool):
        """Guard off means we should not touch Metal limits either."""
        enforcer = ProcessMemoryEnforcer(
            engine_pool=mock_engine_pool,
            memory_guard_tier="balanced",
            prefill_memory_guard=False,
        )
        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch.object(asyncio, "create_task", side_effect=_close_coro),
        ):
            enforcer.start()
        mock_mx.set_wired_limit.assert_not_called()


class TestSingleModelMemoryPressure:
    """Tests for hard-pressure single-model memory handling.

    Hard pressure must reduce resident model memory quickly. Idle models are
    evicted immediately, including the final non-pinned model. Busy models are
    aborted first and kept loaded unless pressure reaches the emergency tier.
    """

    @pytest.mark.asyncio
    async def test_single_idle_model_unloads_at_hard_pressure(self, enforcer):
        """A final idle non-pinned model is unloaded at hard pressure."""
        engine = MagicMock()
        engine.has_active_requests.return_value = False
        engine.abort_all_requests = AsyncMock(return_value=0)
        entry = _make_entry("big-model", engine=engine)
        enforcer._engine_pool._entries = {"big-model": entry}
        enforcer._engine_pool._find_lru_victim.return_value = "big-model"

        async def fake_unload(model_id):
            enforcer._engine_pool._entries[model_id].engine = None

        enforcer._engine_pool._unload_engine.side_effect = fake_unload

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    15 * 1024**3,  # Initial check
                    15 * 1024**3,  # While loop check
                    8 * 1024**3,  # After unload
                ]
            )
            await enforcer._check_and_enforce()

        engine.abort_all_requests.assert_awaited_once()
        enforcer._engine_pool._unload_engine.assert_awaited_once_with("big-model")
        assert entry.engine is None

    @pytest.mark.asyncio
    async def test_single_busy_model_aborts_and_keeps_model_at_hard_pressure(
        self, enforcer
    ):
        """Hard pressure aborts the request first; the final busy model stays loaded."""
        engine = MagicMock()
        engine.has_active_requests.return_value = False
        engine.abort_all_requests = AsyncMock(return_value=3)
        entry = _make_entry("big-model", engine=engine)
        entry.in_use = 1
        enforcer._engine_pool._entries = {"big-model": entry}
        enforcer._engine_pool._find_lru_victim.return_value = None

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    11 * 1024**3,
                    11 * 1024**3,
                ]
            )
            await enforcer._check_and_enforce()

        engine.abort_all_requests.assert_awaited_once()
        enforcer._engine_pool._unload_engine.assert_not_awaited()
        assert entry.engine is not None
        assert entry.pending_unload_reason is None
        assert entry.abort_requested is False

    @pytest.mark.asyncio
    async def test_single_busy_model_pending_unload_at_emergency_pressure(
        self, enforcer
    ):
        """Emergency pressure can still unload the final busy model after drain."""
        engine = MagicMock()
        engine.has_active_requests.return_value = False
        engine.abort_all_requests = AsyncMock(return_value=3)
        entry = _make_entry("big-model", engine=engine)
        entry.in_use = 1
        enforcer._engine_pool._entries = {"big-model": entry}
        enforcer._engine_pool._find_lru_victim.return_value = None

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    13 * 1024**3,
                    13 * 1024**3,
                ]
            )
            await enforcer._check_and_enforce()

        engine.abort_all_requests.assert_awaited_once()
        enforcer._engine_pool._unload_engine.assert_not_awaited()
        assert entry.engine is not None
        assert entry.pending_unload_reason == "hard memory pressure"
        assert entry.abort_requested is True

        entry.in_use = 0
        await enforcer._engine_pool._unload_pending_if_idle_locked("big-model")
        enforcer._engine_pool._unload_engine.assert_awaited_once_with("big-model")

    @pytest.mark.asyncio
    async def test_two_models_one_inferring_evicts_idle(self, enforcer):
        """Scenario 1: Two models, only one inferring. Evict idle LRU."""
        engine_active = MagicMock()
        engine_active.abort_all_requests = AsyncMock(return_value=0)
        engine_idle = MagicMock()
        engine_idle.abort_all_requests = AsyncMock(return_value=0)

        entry_active = _make_entry("active-model", engine=engine_active)
        entry_idle = _make_entry("idle-model", engine=engine_idle)
        enforcer._engine_pool._entries = {
            "active-model": entry_active,
            "idle-model": entry_idle,
        }
        enforcer._engine_pool._find_lru_victim.return_value = "idle-model"

        async def fake_unload(model_id):
            enforcer._engine_pool._entries[model_id].engine = None

        enforcer._engine_pool._unload_engine.side_effect = fake_unload

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.side_effect = _cycling(
                [
                    15 * 1024**3,  # Initial check
                    15 * 1024**3,  # While loop check
                    8 * 1024**3,  # After eviction (under limit)
                ]
            )
            await enforcer._check_and_enforce()

        enforcer._engine_pool._unload_engine.assert_awaited_once_with("idle-model")
        # Idle model's requests aborted before eviction (0 requests)
        engine_idle.abort_all_requests.assert_awaited_once()
        # Active model's requests NOT aborted
        engine_active.abort_all_requests.assert_not_awaited()
        assert entry_active.engine is not None

    @pytest.mark.asyncio
    async def test_two_busy_models_abort_lru_without_pending_at_hard_pressure(
        self, enforcer
    ):
        """Hard pressure aborts the LRU busy model without scheduling unload."""
        engine_a = MagicMock()
        engine_a.has_active_requests.return_value = False
        engine_a.abort_all_requests = AsyncMock(return_value=2)
        engine_b = MagicMock()
        engine_b.has_active_requests.return_value = False
        engine_b.abort_all_requests = AsyncMock(return_value=1)

        entry_a = _make_entry("model-a", engine=engine_a)
        entry_b = _make_entry("model-b", engine=engine_b)
        entry_a.in_use = 1
        entry_b.in_use = 1
        entry_a.last_access = 20
        entry_b.last_access = 10
        enforcer._engine_pool._entries = {
            "model-a": entry_a,
            "model-b": entry_b,
        }
        enforcer._engine_pool._find_lru_victim.return_value = None

        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            # Memory stays over limit throughout
            mock_mx.get_active_memory.return_value = 11 * 1024**3
            await enforcer._check_and_enforce()

        engine_b.abort_all_requests.assert_awaited_once()
        engine_a.abort_all_requests.assert_not_awaited()
        enforcer._engine_pool._unload_engine.assert_not_awaited()
        assert entry_b.pending_unload_reason is None
        assert entry_b.abort_requested is False
        assert entry_a.pending_unload_reason is None


class TestMemoryLimitPropagation:
    """Tests for soft/hard memory limit propagation to schedulers."""

    def test_propagate_memory_limit(self, enforcer):
        """Propagates soft and hard limits to scheduler and batch_generator."""
        bg = MagicMock(spec=[])
        bg._memory_limit_bytes = 0
        bg._memory_hard_limit_bytes = 0
        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = bg
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        enforcer._engine_pool._entries = {"model-a": entry}

        enforcer._propagate_memory_limit()

        # Fixture stubs the ceiling at 10 GB with soft_threshold = 1.0,
        # so the scheduler soft limit equals the ceiling.
        assert scheduler._memory_limit_bytes == 10 * 1024**3
        assert bg._memory_limit_bytes == 10 * 1024**3
        assert scheduler._memory_hard_limit_bytes == 10 * 1024**3
        assert bg._memory_hard_limit_bytes == 10 * 1024**3

    def test_propagate_reserves_hot_cache_used_plus_slack(self, mock_engine_pool):
        """Scheduler limits reserve actual hot-cache bytes plus slack."""
        budget = SimpleNamespace(
            total_bytes=int(5.5 * 1024**3),
            max_bytes=6 * 1024**3,
        )
        mock_engine_pool._scheduler_config = SimpleNamespace(hot_cache_budget=budget)
        enforcer = _make_enforcer(
            mock_engine_pool,
            ceiling=30 * 1024**3,
            soft_threshold=0.92,
        )
        enforcer._get_abort_limit_bytes = lambda: 30 * 1024**3

        bg = MagicMock(spec=[])
        bg._memory_limit_bytes = 0
        bg._memory_hard_limit_bytes = 0
        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = bg
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        mock_engine_pool._entries = {"model-a": _make_entry("model-a", engine=engine)}

        enforcer._propagate_memory_limit()

        expected_hard = 24 * 1024**3
        expected_soft = int(expected_hard * 0.92)
        assert scheduler._memory_hard_limit_bytes == expected_hard
        assert scheduler._memory_limit_bytes == expected_soft
        assert scheduler._memory_abort_limit_bytes == expected_hard
        assert bg._memory_hard_limit_bytes == expected_hard
        assert bg._memory_limit_bytes == expected_soft

    def test_propagate_charges_only_slack_when_hot_cache_empty(self, mock_engine_pool):
        budget = SimpleNamespace(total_bytes=0, max_bytes=6 * 1024**3)
        mock_engine_pool._scheduler_config = SimpleNamespace(hot_cache_budget=budget)
        enforcer = _make_enforcer(mock_engine_pool, ceiling=30 * 1024**3)
        enforcer._get_abort_limit_bytes = lambda: 30 * 1024**3

        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = None
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        mock_engine_pool._entries = {"model-a": _make_entry("model-a", engine=engine)}

        enforcer._propagate_memory_limit()

        assert scheduler._memory_hard_limit_bytes == int(29.5 * 1024**3)

    def test_propagate_scheduler_limit_clamps_to_one_byte(self, mock_engine_pool):
        budget = SimpleNamespace(total_bytes=80 * 1024**3, max_bytes=80 * 1024**3)
        mock_engine_pool._scheduler_config = SimpleNamespace(hot_cache_budget=budget)
        enforcer = _make_enforcer(mock_engine_pool, ceiling=30 * 1024**3)
        enforcer._get_abort_limit_bytes = lambda: 30 * 1024**3

        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = None
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        mock_engine_pool._entries = {"model-a": _make_entry("model-a", engine=engine)}

        enforcer._propagate_memory_limit()

        assert scheduler._memory_hard_limit_bytes == 1
        assert scheduler._memory_abort_limit_bytes == 1

    def test_propagate_with_guard_disabled(self, enforcer):
        """When the guard is disabled the field reflects it; hard limit is
        still propagated for observability — the reader's early-return on
        ``prefill_memory_guard=False`` makes the value moot for the
        rejection path."""
        scheduler = MagicMock(spec=[])
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        enforcer._engine_pool._entries = {"model-a": entry}
        enforcer._prefill_memory_guard = False

        enforcer._propagate_memory_limit()

        assert scheduler._prefill_memory_guard is False
        # Hard limit is still propagated for observability — the reader's
        # early-return on ``prefill_memory_guard=False`` makes the value
        # moot for the rejection path. (The fixture's monkey-patched
        # ceiling stays at 10 GB; the production ``_get_hard_limit_bytes``
        # would return 0 when the guard is disabled, but the fixture
        # bypasses that branch — see ``_make_enforcer``.)
        assert scheduler._memory_hard_limit_bytes == 10 * 1024**3

    def test_propagates_on_tier_change(self, enforcer):
        """Changing the tier at runtime triggers re-propagation."""
        bg = MagicMock(spec=[])
        bg._memory_limit_bytes = 0
        bg._memory_hard_limit_bytes = 0
        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = bg
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        enforcer._engine_pool._entries = {"model-a": entry}

        enforcer._running = True
        # Simulate the ceiling shrinking after the tier flip. The new
        # propagation path reads from ``_get_ceiling_breakdown``, not
        # ``_get_hard_limit_bytes`` — patch both for completeness.
        new_ceiling = 20 * 1024**3
        enforcer._get_hard_limit_bytes = lambda: new_ceiling
        enforcer._get_ceiling_breakdown = lambda: {
            "static": new_ceiling,
            "dynamic": new_ceiling,
            "metal_cap": new_ceiling,
            "hard_limit": new_ceiling,
        }
        enforcer.memory_guard_tier = "safe"

        assert scheduler._memory_limit_bytes == new_ceiling
        assert bg._memory_limit_bytes == new_ceiling

    def test_skips_engine_without_scheduler(self, enforcer):
        """Gracefully skips engines without scheduler attribute."""
        engine = MagicMock(spec=[])
        # No scheduler attribute (spec=[] prevents auto-creation)
        entry = _make_entry("model-a", engine=engine)
        enforcer._engine_pool._entries = {"model-a": entry}

        # Should not raise
        enforcer._propagate_memory_limit()

    def test_propagates_to_multiple_engines(self, enforcer):
        """Propagates to all engines."""
        schedulers = []
        entries = {}
        for i in range(3):
            bg = MagicMock(spec=[])
            bg._memory_limit_bytes = 0
            scheduler = MagicMock(spec=[])
            scheduler.batch_generator = bg
            schedulers.append(scheduler)
            engine = MagicMock(spec=[])
            engine.scheduler = scheduler
            entry = _make_entry(f"model-{i}", engine=engine)
            entries[f"model-{i}"] = entry
        enforcer._engine_pool._entries = entries

        enforcer._propagate_memory_limit()

        for scheduler in schedulers:
            assert scheduler._memory_limit_bytes == 10 * 1024**3

    async def test_check_and_enforce_propagates_every_poll(self, enforcer):
        """Regression: a fresh engine loaded AFTER enforcer.start() must pick
        up its limits within one poll interval — even when pressure stays
        "ok" the whole time.

        Before this guarantee the propagation only fired on pressure-level
        changes. On a host where the first prefill stayed below soft until
        a few seconds in, the scheduler kept _prefill_memory_guard=False /
        _memory_hard_limit_bytes=0 (their __init__ defaults), the guard
        short-circuited, the request entered prefill, and the underlying
        Apple IOGPUFamily bug (FB22091885) panicked the kernel mid-chunk.
        """
        # Engine pool starts empty (mirrors real startup: lazy load on first
        # request, well after enforcer.start()).
        enforcer._engine_pool._entries = {}
        # Engine loads at t1 — the enforcer hasn't seen it yet.
        bg = MagicMock(spec=[])
        bg._memory_limit_bytes = 0
        bg._memory_hard_limit_bytes = 0
        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = bg
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        enforcer._engine_pool._entries = {"model-a": entry}

        # One poll iteration with pressure well below soft — pressure level
        # does NOT change. Before the fix this returned without propagating.
        with patch.object(enforcer, "_current_usage_bytes", return_value=1 * 1024**3):
            await enforcer._check_and_enforce()

        # Within one poll, the freshly-loaded engine has the user-configured
        # ceiling and the guard flag.
        assert scheduler._memory_hard_limit_bytes == 10 * 1024**3
        assert scheduler._memory_limit_bytes == 10 * 1024**3
        assert scheduler._prefill_memory_guard is True

    def test_propagates_through_batched_engine_wrapper(self, enforcer):
        """Regression: live engines in EnginePool don't expose ``.scheduler``
        on the top-level wrapper — BatchedEngine and VLMBatchedEngine both
        hold the real Scheduler at ``self._engine.engine.scheduler``. The
        propagation must traverse that chain, otherwise the prefill memory
        guard flag never reaches the scheduler and the guard short-circuits
        on every request (observed end-to-end 2026-05-15: three kernel
        panics from 110k-token Qwen3.6-VL prefills the guard "should" have
        rejected).
        """
        # Build the real wrapper shape:
        #   entry.engine                  → BatchedEngine / VLMBatchedEngine
        #   entry.engine._engine          → AsyncEngineCore
        #   entry.engine._engine.engine   → EngineCore
        #   entry.engine._engine.engine.scheduler → Scheduler  ← target
        scheduler = MagicMock(spec=[])
        scheduler.batch_generator = None
        engine_core = MagicMock(spec=["scheduler"])
        engine_core.scheduler = scheduler
        async_engine_core = MagicMock(spec=["engine"])
        async_engine_core.engine = engine_core
        # Wrapper deliberately does NOT expose top-level ``.scheduler`` — only
        # ``._engine`` like the real BatchedEngine.
        wrapper = MagicMock(spec=["_engine"])
        wrapper._engine = async_engine_core

        entry = _make_entry("model-a", engine=wrapper)
        enforcer._engine_pool._entries = {"model-a": entry}

        enforcer._propagate_memory_limit()

        assert scheduler._memory_limit_bytes == 10 * 1024**3
        assert scheduler._memory_hard_limit_bytes == 10 * 1024**3
        assert scheduler._prefill_memory_guard is True

    def test_unresolvable_scheduler_logs_warning_once(self, enforcer, caplog):
        """If the wrapper-chain traversal fails (no ``scheduler`` anywhere
        in the chain), ``_propagate_memory_limit`` must log a WARNING
        naming the engine type so the silent no-op failure mode that
        originally hid the dead memory guard is loud in CI / oncall. The
        warning is rate-limited per engine type so a misconfigured
        engine polled every second doesn't spam.
        """
        # Wrapper chain that bottoms out without a scheduler.
        wrapper = MagicMock(spec=["_engine"])
        wrapper._engine = MagicMock(spec=["engine"])
        wrapper._engine.engine = MagicMock(spec=[])  # no .scheduler
        wrapper.__class__.__name__ = "BrokenEngine"

        entry = _make_entry("model-broken", engine=wrapper)
        enforcer._engine_pool._entries = {"model-broken": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()
            # Second call: no extra log line — rate limit holds.
            enforcer._propagate_memory_limit()

        matching = [
            r for r in caplog.records if "could not resolve scheduler" in r.message
        ]
        assert (
            len(matching) == 1
        ), f"expected 1 warning, got {[r.message for r in matching]}"


class TestUnresolvableSchedulerWarning:
    """``_propagate_memory_limit`` must complain when a wrapper-chain
    change makes the scheduler unreachable via ``_resolve_scheduler``.

    Silent no-op was the failure mode that originally hid the
    dead-memory-guard bug for months — surfacing it as a per-engine-type
    rate-limited warning closes that gap without spamming logs at
    request rate.
    """

    def test_warns_once_per_engine_type(self, enforcer, caplog):
        """Three propagation calls with an unreachable scheduler emit
        exactly one WARNING (not three)."""
        # Engine wrapper with no resolvable scheduler chain: spec=[]
        # blocks both ``.scheduler`` and ``._engine`` lookups so
        # ``_resolve_scheduler`` returns None.
        engine = MagicMock(spec=[])
        entry = _make_entry("model-broken", engine=engine)
        enforcer._engine_pool._entries = {"model-broken": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()
            enforcer._propagate_memory_limit()
            enforcer._propagate_memory_limit()

        warnings = [
            r for r in caplog.records if "could not resolve scheduler" in r.getMessage()
        ]
        assert len(warnings) == 1, (
            f"Expected exactly one warning per engine type per lifetime, "
            f"got {len(warnings)}"
        )

    def test_diffusion_engine_without_scheduler_does_not_warn(self, enforcer, caplog):
        """Diffusion VLMs intentionally bypass AsyncEngineCore schedulers."""
        engine = MagicMock(spec=["is_diffusion_model"])
        engine.is_diffusion_model = True
        entry = _make_entry("model-diffusion", engine=engine)
        enforcer._engine_pool._entries = {"model-diffusion": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()

        warnings = [
            r for r in caplog.records if "could not resolve scheduler" in r.getMessage()
        ]
        assert warnings == []

    def test_unresolvable_does_not_block_other_engines(self, enforcer):
        """If engine A is unresolvable but engine B has a real scheduler,
        B must still receive the propagation."""
        # A: unresolvable (no .scheduler / no ._engine).
        engine_a = MagicMock(spec=[])
        entry_a = _make_entry("model-a", engine=engine_a)

        # B: real scheduler chain.
        scheduler_b = MagicMock(spec=[])
        scheduler_b._memory_limit_bytes = 0
        scheduler_b._memory_hard_limit_bytes = 0
        scheduler_b._prefill_memory_guard = False
        scheduler_b._admission_paused = False
        scheduler_b._prefill_safe_zone_ratio = 0.0
        scheduler_b._prefill_min_chunk_tokens = 0
        scheduler_b.batch_generator = None
        engine_b = MagicMock(spec=[])
        engine_b.scheduler = scheduler_b
        entry_b = _make_entry("model-b", engine=engine_b)

        enforcer._engine_pool._entries = {
            "model-a": entry_a,
            "model-b": entry_b,
        }

        enforcer._propagate_memory_limit()

        # B got it; A's resolve-failure didn't poison the loop.
        assert scheduler_b._memory_limit_bytes == 10 * 1024**3
        assert scheduler_b._prefill_memory_guard is True

    def test_no_warning_for_unloaded_engine(self, enforcer, caplog):
        """A discovered-but-unloaded entry (``engine is None``) is a normal
        state, not a wrapper break, so it must not emit the warning. The
        pool keeps these entries for every model it has discovered but not
        yet loaded, so warning here would fire on a routine startup."""
        entry = _make_entry("model-unloaded", engine=None)
        enforcer._engine_pool._entries = {"model-unloaded": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()
            enforcer._propagate_memory_limit()

        warnings = [
            r for r in caplog.records if "could not resolve scheduler" in r.getMessage()
        ]
        assert (
            warnings == []
        ), f"Unloaded engine must not warn, got {len(warnings)} warning(s)"

    def test_no_warning_for_stopped_engine_during_unload(self, enforcer, caplog):
        """During EnginePool unload there is a short window where the engine
        object still sits on the entry but its scheduler has already been
        released. That is teardown, not a broken wrapper chain.
        """
        engine = MagicMock(spec=["_loaded"])
        engine._loaded = False
        entry = _make_entry("model-stopped", engine=engine)
        enforcer._engine_pool._entries = {"model-stopped": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()

        warnings = [
            r for r in caplog.records if "could not resolve scheduler" in r.getMessage()
        ]
        assert warnings == []

    def test_no_warning_for_dflash_without_fallback(self, enforcer, caplog):
        """DFlash normal mode has no scheduler until fallback is started."""

        class DFlashEngine:
            _fallback_engine = None

        entry = _make_entry("model-dflash", engine=DFlashEngine())
        enforcer._engine_pool._entries = {"model-dflash": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()

        warnings = [
            r for r in caplog.records if "could not resolve scheduler" in r.getMessage()
        ]
        assert warnings == []


class TestStoreCacheCapWalk:
    """Tests for _walk_store_cache_caps — store-cache gate adjustment (#1383)."""

    def _scheduler_with_adjust(self):
        scheduler = MagicMock(spec=[])
        scheduler.adjust_store_cache_cap = MagicMock()
        return scheduler

    def test_calls_adjust_with_current_pressure(self, enforcer):
        scheduler = self._scheduler_with_adjust()
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        enforcer._engine_pool._entries = {"m": _make_entry("m", engine=engine)}
        enforcer._pressure_level = "soft"

        enforcer._walk_store_cache_caps()

        scheduler.adjust_store_cache_cap.assert_called_once_with("soft")

    def test_no_op_when_engine_missing(self, enforcer):
        entry = _make_entry("m", engine=None)
        enforcer._engine_pool._entries = {"m": entry}
        # Should not raise.
        enforcer._walk_store_cache_caps()

    def test_no_op_when_scheduler_lacks_method(self, enforcer):
        engine = MagicMock(spec=[])  # no scheduler attr
        entry = _make_entry("m", engine=engine)
        enforcer._engine_pool._entries = {"m": entry}
        # Should not raise.
        enforcer._walk_store_cache_caps()

    @pytest.mark.asyncio
    async def test_check_and_enforce_walks_caps_on_ok(self, enforcer):
        scheduler = self._scheduler_with_adjust()
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        enforcer._engine_pool._entries = {"m": _make_entry("m", engine=engine)}

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint", return_value=0),
        ):
            mock_mx.get_active_memory.return_value = 1 * 1024**3  # ok
            await enforcer._check_and_enforce()

        scheduler.adjust_store_cache_cap.assert_called_with("ok")

    @pytest.mark.asyncio
    async def test_check_and_enforce_walks_caps_on_soft(self, enforcer):
        # Force a 0.85/0.95 split so 9GB lands in the soft band.
        enforcer._soft_threshold = 0.85
        enforcer._hard_threshold = 0.95
        scheduler = self._scheduler_with_adjust()
        engine = MagicMock(spec=[])
        engine.scheduler = scheduler
        enforcer._engine_pool._entries = {"m": _make_entry("m", engine=engine)}
        enforcer._engine_pool._find_lru_victim = MagicMock(return_value=None)

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint", return_value=0),
        ):
            mock_mx.get_active_memory.return_value = 9 * 1024**3  # soft
            await enforcer._check_and_enforce()

        scheduler.adjust_store_cache_cap.assert_called_with("soft")


class TestProperties:
    """Tests for enforcer properties."""

    def test_memory_guard_tier_default(self, enforcer):
        """Default tier from `_make_enforcer` is balanced."""
        assert enforcer.memory_guard_tier == "balanced"

    def test_memory_guard_tier_setter(self, enforcer):
        """Setting a new tier updates the internal state."""
        enforcer.memory_guard_tier = "safe"
        assert enforcer.memory_guard_tier == "safe"

    def test_memory_guard_tier_setter_ignores_unknown_value(self, enforcer):
        """Unknown tier values normalize to balanced."""
        enforcer.memory_guard_tier = "extreme"
        assert enforcer.memory_guard_tier == "balanced"

    def test_get_final_ceiling_matches_hard_limit(self, enforcer):
        assert enforcer.get_final_ceiling() == enforcer._get_hard_limit_bytes()

    def test_is_running_initially_false(self, enforcer):
        """Test is_running is False before start."""
        assert enforcer.is_running is False

    def test_get_status_when_not_running(self, enforcer):
        """Test get_status when enforcer is not running."""
        status = enforcer.get_status()
        assert status["enabled"] is False
        assert status["ceiling_bytes"] == 0
        assert status["current_bytes"] == 0
        assert status["memory_guard_tier"] == "balanced"


class TestLifecycle:
    """Tests for start/stop lifecycle."""

    @pytest.mark.asyncio
    async def test_start_stop(self, enforcer):
        """Test start and stop lifecycle."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 0
            enforcer.start()
            assert enforcer.is_running is True
            await asyncio.sleep(0.05)
            await enforcer.stop()
            assert enforcer.is_running is False

    @pytest.mark.asyncio
    async def test_double_start_is_noop(self, enforcer):
        """Test calling start twice doesn't create duplicate tasks."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 0
            enforcer.start()
            task1 = enforcer._task
            enforcer.start()
            task2 = enforcer._task
            assert task1 is task2
            await enforcer.stop()

    @pytest.mark.asyncio
    async def test_stop_when_not_started(self, enforcer):
        """Test stop when not started is safe."""
        await enforcer.stop()  # Should not raise

    @pytest.mark.asyncio
    async def test_get_status_when_running(self, enforcer):
        """Test get_status reflects running state."""
        with patch("omlx.process_memory_enforcer.mx") as mock_mx:
            mock_mx.get_active_memory.return_value = 5 * 1024**3
            enforcer.start()
            status = enforcer.get_status()
            assert status["enabled"] is True
            assert status["current_bytes"] == 5 * 1024**3
            await enforcer.stop()


class TestTwoWatermarkPressureLevels:
    """Tests for 2-watermark soft/hard pressure level handling."""

    @pytest.fixture
    def pool(self):
        p = MagicMock()
        p._lock = asyncio.Lock()
        p._find_lru_victim = MagicMock(return_value=None)
        p._unload_engine = AsyncMock()
        p._find_pending_unload_ready_locked = MagicMock(return_value=None)
        p._unload_pending_if_idle_locked = AsyncMock(return_value=False)
        p._mark_pending_unload_locked = MagicMock(return_value=False)
        p._entries = {}
        return p

    @pytest.fixture
    def enforcer_2wm(self, pool):
        return _make_enforcer(
            pool,
            ceiling=100 * 1024**3,
            soft_threshold=0.85,
            hard_threshold=0.95,
        )

    def test_soft_hard_bytes_computed(self, enforcer_2wm):
        assert enforcer_2wm._soft_bytes() == int(100 * 1024**3 * 0.85)
        assert enforcer_2wm._hard_bytes() == int(100 * 1024**3 * 0.95)

    def test_prefill_abort_margin_is_tier_specific(self, pool):
        balanced = _make_enforcer(pool, tier="balanced")
        aggressive = _make_enforcer(pool, tier="aggressive")
        custom = _make_enforcer(pool, tier="custom")

        assert balanced._get_prefill_abort_margin() == 0.90
        assert aggressive._get_prefill_abort_margin() == 0.95
        assert custom._get_prefill_abort_margin() == 0.95

    @pytest.mark.parametrize(
        "tier,expected",
        [("safe", 0.85), ("balanced", 0.90), ("aggressive", 0.925)],
    )
    def test_legacy_default_soft_threshold_is_tier_specific(self, pool, tier, expected):
        enforcer = ProcessMemoryEnforcer(
            pool,
            memory_guard_tier=tier,
            soft_threshold=0.85,
        )
        assert enforcer._soft_threshold == expected

    def test_explicit_soft_threshold_override_is_preserved(self, pool):
        enforcer = ProcessMemoryEnforcer(
            pool,
            memory_guard_tier="aggressive",
            soft_threshold=0.92,
        )
        assert enforcer._soft_threshold == 0.92

    def test_tier_change_refreshes_soft_threshold_and_prefill_headroom(self, pool):
        enforcer = ProcessMemoryEnforcer(pool, memory_guard_tier="safe")

        assert enforcer._soft_threshold == 0.85
        assert enforcer._get_prefill_headroom_safety() == 0.90

        enforcer.memory_guard_tier = "aggressive"

        assert enforcer._soft_threshold == 0.925
        assert enforcer._get_prefill_headroom_safety() == 0.925

    def test_prefill_headroom_safety_is_tier_specific(self, pool):
        balanced = ProcessMemoryEnforcer(pool, memory_guard_tier="balanced")
        aggressive = ProcessMemoryEnforcer(pool, memory_guard_tier="aggressive")

        assert balanced._get_prefill_headroom_safety() == 0.90
        assert aggressive._get_prefill_headroom_safety() == 0.925

    def test_get_pressure_level_when_not_running(self, enforcer_2wm):
        # _running=False → always ok regardless of cached level
        enforcer_2wm._pressure_level = "hard"
        assert enforcer_2wm.get_pressure_level() == "ok"

    def test_get_pressure_level_when_running_returns_cached(self, enforcer_2wm):
        enforcer_2wm._running = True
        enforcer_2wm._pressure_level = "soft"
        assert enforcer_2wm.get_pressure_level() == "soft"

    @pytest.mark.asyncio
    async def test_ok_when_below_soft(self, enforcer_2wm):
        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 50 * 1024**3
            await enforcer_2wm._check_and_enforce()
        assert enforcer_2wm._pressure_level == "ok"
        enforcer_2wm._engine_pool._unload_engine.assert_not_called()

    @pytest.mark.asyncio
    async def test_soft_when_active_low_but_phys_high(self, enforcer_2wm):
        """phys_footprint dominates active — the #702 case."""
        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            # active well below soft, phys above soft but below hard
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 88 * 1024**3
            await enforcer_2wm._check_and_enforce()
        assert enforcer_2wm._pressure_level == "soft"

    @pytest.mark.asyncio
    async def test_hard_when_phys_at_hard_threshold(self, enforcer_2wm):
        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 60 * 1024**3
            gpf.return_value = 98 * 1024**3
            await enforcer_2wm._check_and_enforce()
        assert enforcer_2wm._pressure_level == "hard"

    @pytest.mark.asyncio
    async def test_hard_pressure_shrinks_hot_cache_before_abort(self, mock_engine_pool):
        budget = MagicMock()
        budget.total_bytes = 20 * 1024**3
        budget.max_bytes = 20 * 1024**3
        budget.shrink_to.return_value = 8 * 1024**3
        mock_engine_pool._scheduler_config = SimpleNamespace(hot_cache_budget=budget)
        enforcer = _make_enforcer(
            mock_engine_pool,
            ceiling=100 * 1024**3,
            soft_threshold=0.90,
            hard_threshold=0.95,
        )
        to_thread_calls = []

        async def run_inline(fn, *args, **kwargs):
            to_thread_calls.append((fn, args, kwargs))
            return fn(*args, **kwargs)

        with (
            patch.object(
                enforcer,
                "_current_usage_bytes",
                side_effect=[98 * 1024**3, 80 * 1024**3],
            ),
            patch(
                "omlx.process_memory_enforcer.asyncio.to_thread",
                side_effect=run_inline,
            ),
        ):
            await enforcer._check_and_enforce()

        assert to_thread_calls
        assert to_thread_calls[0][0] == enforcer._shrink_hot_cache_for_pressure
        budget.shrink_to.assert_called_once()
        target_hot = budget.shrink_to.call_args.args[0]
        assert target_hot == 12 * 1024**3
        assert enforcer._pressure_level == "ok"
        mock_engine_pool._unload_engine.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_propagates_admission_paused_on_soft(self, enforcer_2wm, pool):
        # Wire a scheduler-like mock so propagate has something to set.
        engine = MagicMock()
        scheduler = MagicMock()
        engine.scheduler = scheduler
        entry = _make_entry("m", engine=engine)
        pool._entries = {"m": entry}

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 88 * 1024**3
            await enforcer_2wm._check_and_enforce()

        assert scheduler._admission_paused is True

    @pytest.mark.asyncio
    async def test_clears_admission_paused_on_recovery(self, enforcer_2wm, pool):
        engine = MagicMock()
        scheduler = MagicMock()
        scheduler._admission_paused = True
        engine.scheduler = scheduler
        entry = _make_entry("m", engine=engine, is_pinned=True)
        pool._entries = {"m": entry}

        # Force into soft first
        enforcer_2wm._pressure_level = "soft"

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 30 * 1024**3
            gpf.return_value = 40 * 1024**3
            await enforcer_2wm._check_and_enforce()

        assert enforcer_2wm._pressure_level == "ok"
        assert scheduler._admission_paused is False

    @pytest.mark.asyncio
    async def test_propagates_abort_limit_to_scheduler(self, enforcer_2wm, pool):
        """The stable abort ceiling is pushed to scheduler._memory_abort_limit_bytes
        every tick, independent of the (jittery) dynamic hard limit."""
        engine = MagicMock()
        scheduler = MagicMock()
        scheduler._memory_limit_bytes = 0
        scheduler._memory_hard_limit_bytes = 0
        scheduler._memory_abort_limit_bytes = 0
        scheduler._prefill_memory_guard = False
        scheduler._admission_paused = False
        engine.scheduler = scheduler
        pool._entries = {"m": _make_entry("m", engine=engine)}

        # Stub the abort ceiling to a known stable value.
        enforcer_2wm._get_abort_limit_bytes = lambda: 42 * 1024**3

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 50 * 1024**3
            await enforcer_2wm._check_and_enforce()

        assert scheduler._memory_abort_limit_bytes == 42 * 1024**3

    @pytest.mark.asyncio
    async def test_propagates_prefill_abort_margin_to_scheduler(self, pool):
        enforcer = _make_enforcer(pool, ceiling=100 * 1024**3, tier="custom")
        engine = MagicMock()
        scheduler = MagicMock()
        scheduler._memory_limit_bytes = 0
        scheduler._memory_hard_limit_bytes = 0
        scheduler._memory_abort_limit_bytes = 0
        scheduler._prefill_abort_margin = 0.90
        scheduler._prefill_memory_guard = False
        scheduler._admission_paused = False
        engine.scheduler = scheduler
        pool._entries = {"m": _make_entry("m", engine=engine)}

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 50 * 1024**3
            await enforcer._check_and_enforce()

        assert scheduler._prefill_abort_margin == 0.95

    @pytest.mark.asyncio
    async def test_hard_below_ceiling_does_not_abort_all_pinned(
        self, enforcer_2wm, pool
    ):
        engine = MagicMock()
        engine.abort_all_requests = AsyncMock(return_value=3)
        entry = _make_entry("pinned", engine=engine, is_pinned=True)
        pool._entries = {"pinned": entry}
        # Single pinned model means find_lru_victim returns None (pinned not victim).
        pool._find_lru_victim.return_value = None

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 60 * 1024**3
            gpf.return_value = 99 * 1024**3
            await enforcer_2wm._check_and_enforce()

        # No in-progress loads to abort, all pinned → enforcer just logs warning,
        # doesn't crash.
        assert enforcer_2wm._pressure_level == "hard"
        engine.abort_all_requests.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_emergency_margin_aborts_all_pinned_requests(
        self, enforcer_2wm, pool
    ):
        engine = MagicMock()
        engine.abort_all_requests = AsyncMock(return_value=3)
        entry = _make_entry("pinned", engine=engine, is_pinned=True)
        pool._entries = {"pinned": entry}
        pool._find_lru_victim.return_value = None

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 60 * 1024**3
            gpf.return_value = 103 * 1024**3
            await enforcer_2wm._check_and_enforce()

        engine.abort_all_requests.assert_awaited_once()
        pool._unload_engine.assert_not_awaited()
        assert entry.engine is engine

    @pytest.mark.asyncio
    async def test_emergency_consecutive_over_ceiling_polls_abort_all_pinned(
        self, enforcer_2wm, pool
    ):
        engine = MagicMock()
        engine.abort_all_requests = AsyncMock(return_value=2)
        entry = _make_entry("pinned", engine=engine, is_pinned=True)
        pool._entries = {"pinned": entry}
        pool._find_lru_victim.return_value = None

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 60 * 1024**3
            gpf.return_value = 101 * 1024**3
            await enforcer_2wm._check_and_enforce()

        engine.abort_all_requests.assert_not_awaited()
        engine.abort_all_requests.reset_mock()

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 60 * 1024**3
            gpf.return_value = 101 * 1024**3
            await enforcer_2wm._check_and_enforce()

        engine.abort_all_requests.assert_awaited_once()
        pool._unload_engine.assert_not_awaited()
        assert entry.engine is engine

    @pytest.mark.asyncio
    async def test_soft_does_not_abort_loading(self, enforcer_2wm, pool):
        loading_entry = _make_entry("loading", engine=None, is_loading=True)
        pool._entries = {"loading": loading_entry}
        pool._find_lru_victim.return_value = None

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 88 * 1024**3  # soft
            await enforcer_2wm._check_and_enforce()

        assert loading_entry.abort_loading is False  # soft must not abort load

    @pytest.mark.asyncio
    async def test_hard_aborts_loading(self, enforcer_2wm, pool):
        loading_entry = _make_entry("loading", engine=None, is_loading=True)
        pool._entries = {"loading": loading_entry}
        pool._find_lru_victim.return_value = None

        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 60 * 1024**3
            gpf.return_value = 99 * 1024**3  # hard
            await enforcer_2wm._check_and_enforce()

        assert loading_entry.abort_loading is True

    def test_get_status_uses_max_active_and_phys(self, enforcer_2wm):
        """get_status must report the same value enforcer compares against,
        so admin UI / /health utilization matches the watermark logic."""
        enforcer_2wm._running = True
        with (
            patch("omlx.process_memory_enforcer.mx") as mock_mx,
            patch("omlx.process_memory_enforcer.get_phys_footprint") as gpf,
        ):
            mock_mx.get_active_memory.return_value = 50 * 1024**3
            gpf.return_value = 88 * 1024**3  # phys dominates
            status = enforcer_2wm.get_status()
        assert status["current_bytes"] == 88 * 1024**3
        # Utilization computed against the max value
        assert abs(status["utilization"] - 0.88) < 0.01


class TestDFlashGuardPropagation:
    """The enforcer must reach DFlash's primary-mode guard target.

    DFlash bypasses the scheduler: in primary mode it exposes a lightweight
    ``_prefill_guard``; in fallback mode its ``scheduler`` property resolves
    the fallback engine's real scheduler (covered by test_dflash_engine.py).
    Without the ``_prefill_guard`` arm in ``_resolve_scheduler`` the watermarks
    never reach a primary-mode DFlash and the prefill guard stays dead.
    """

    def test_resolves_primary_guard(self, enforcer):
        guard = MagicMock(spec=[])
        engine = MagicMock(spec=["_prefill_guard"])
        engine._prefill_guard = guard
        entry = _make_entry("dflash", engine=engine)
        enforcer._engine_pool._entries = {"dflash": entry}

        assert enforcer._resolve_scheduler(entry) is guard

        enforcer._propagate_memory_limit()
        assert guard._memory_hard_limit_bytes == 10 * 1024**3
        assert guard._prefill_memory_guard == enforcer._prefill_memory_guard

    def test_dflash_primary_does_not_warn(self, enforcer, caplog):
        engine = MagicMock(spec=["_prefill_guard"])
        engine._prefill_guard = MagicMock(spec=[])
        entry = _make_entry("dflash", engine=engine)
        enforcer._engine_pool._entries = {"dflash": entry}

        with caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"):
            enforcer._propagate_memory_limit()

        assert not [
            r for r in caplog.records if "could not resolve scheduler" in r.message
        ]

    def test_non_dflash_resolution_unchanged(self, enforcer):
        """A standard engine with a direct ``.scheduler`` still resolves via the
        existing chain — the DFlash arm must not interfere."""
        scheduler = MagicMock(spec=[])
        engine = MagicMock(spec=["scheduler"])
        engine.scheduler = scheduler
        entry = _make_entry("model-a", engine=engine)
        assert enforcer._resolve_scheduler(entry) is scheduler


class TestWiredLimitSuggestionClamp:
    """The recommended iogpu.wired_limit_mb must leave the OS headroom (#2184)."""

    def _with_total(self, total):
        return patch("omlx.settings.get_system_memory", return_value=total)

    def test_suggestion_below_reserve_unchanged(self):
        total = 512 * 1024**3
        desired = 400 * 1024**3
        with self._with_total(total):
            assert pme._wired_limit_suggestion_bytes(desired) == desired

    def test_near_physical_suggestion_clamped(self):
        total = 512 * 1024**3
        desired = total - 8 * 1024**3  # safe-tier static ceiling: 504 GiB
        with self._with_total(total):
            clamped = pme._wired_limit_suggestion_bytes(desired)
        # 5% of 512 GiB = 25.6 GiB reserve
        assert clamped == total - total // 20
        assert clamped < desired

    def test_small_mac_recommendation_unchanged(self):
        """Small-memory Macs keep the tier recommendation: 5% of RAM stays
        below the tier static reserve there, so the clamp does not bite and
        users can wire as much as before (#2184 targets large boxes)."""
        total = 32 * 1024**3
        desired = total - 2 * 1024**3  # custom tier ceiling: RAM - 2 GiB
        with self._with_total(total):
            assert pme._wired_limit_suggestion_bytes(desired) == desired

    def test_desired_above_cap_clamps_to_cap(self):
        total = 4 * 1024**3
        with self._with_total(total):
            clamped = pme._wired_limit_suggestion_bytes(8 * 1024**3)
        assert clamped == total - total // 20

    def test_log_hint_uses_clamped_value(self, caplog):
        total = 512 * 1024**3
        desired = total - 8 * 1024**3
        with (
            self._with_total(total),
            patch.object(pme, "get_iogpu_wired_limit_bytes", return_value=0),
            patch.object(
                pme,
                "_get_max_metal_working_set_bytes",
                return_value=384 * 1024**3,
            ),
            caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"),
        ):
            pme._apply_metal_wired_limit(desired)
        hints = [r for r in caplog.records if "iogpu.wired_limit_mb=" in r.message]
        assert hints
        suggested_mb = (total - total // 20) // (1024**2)
        assert f"iogpu.wired_limit_mb={suggested_mb}" in hints[0].getMessage()

    def test_reasonable_user_cap_not_nagged(self, caplog):
        """A user cap at the recommended level must not trigger the raise hint
        even when the static ceiling is higher (the #2184 report followed the
        old hint into a jetsam crash-loop)."""
        total = 512 * 1024**3
        desired = total - 8 * 1024**3  # 504 GiB ceiling
        user_cap = total - total // 20  # exactly the recommendation
        with (
            self._with_total(total),
            patch.object(
                pme, "get_iogpu_wired_limit_bytes", return_value=user_cap
            ),
            patch.object(pme.mx, "set_wired_limit", return_value=0),
            caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"),
        ):
            pme._apply_metal_wired_limit(desired)
        assert not [
            r for r in caplog.records if "Raise it with" in r.message
        ]

    def test_near_physical_user_cap_warns_jetsam(self, caplog):
        total = 512 * 1024**3
        user_cap = total - 2 * 1024**3  # 510 GiB, the crash-loop setting
        with (
            self._with_total(total),
            patch.object(
                pme, "get_iogpu_wired_limit_bytes", return_value=user_cap
            ),
            patch.object(pme.mx, "set_wired_limit", return_value=0),
            caplog.at_level("WARNING", logger="omlx.process_memory_enforcer"),
        ):
            pme._apply_metal_wired_limit(400 * 1024**3)
        assert [r for r in caplog.records if "jetsam" in r.message]