box/tests/test_box_stability_watcher.py

#!/usr/bin/env python3
"""test_box_stability_watcher.py — Comprehensive unit tests for Box Stability Watcher."""

import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from unittest import mock

REPO_ROOT = Path("/home/super/Projects/NetVM")
WATCHERS_DIR = REPO_ROOT / "watchers"
sys.path.insert(0, str(WATCHERS_DIR))

import importlib.util
spec = importlib.util.spec_from_file_location("box_stability_watcher", str(WATCHERS_DIR / "box-stability-watcher.py"))
w = importlib.util.module_from_spec(spec)
spec.loader.exec_module(w)


class TestConfigAndSafety(unittest.TestCase):
    def test_load_config_defaults(self):
        with tempfile.TemporaryDirectory() as td:
            non_existent = Path(td) / "missing.json"
            cfg = w.load_config(non_existent)
            self.assertIn("thresholds", cfg)
            self.assertIn("protected_commands", cfg)
            self.assertEqual(cfg["thresholds"]["load_warning"], 20.0)

    def test_is_protected(self):
        cfg = {"protected_commands": ["sshd", "tailscaled", "tmux", "systemd", "ghostty"]}
        self.assertTrue(w.is_protected(1, "systemd", "/sbin/init", cfg))
        self.assertTrue(w.is_protected(os.getpid(), "python3", "some_script", cfg))
        self.assertTrue(w.is_protected(999, "sshd", "/usr/sbin/sshd -D", cfg))
        self.assertTrue(w.is_protected(888, "tmux", "tmux new-session -s main", cfg))
        self.assertTrue(w.is_protected(777, "tailscaled", "/usr/sbin/tailscaled", cfg))
        self.assertFalse(w.is_protected(1234, "muse-bin", "/home/super/.local/bin/muse-bin-1.4.3 resume abc", cfg))
        self.assertFalse(w.is_protected(5678, "chromium", "/usr/lib/chromium/chromium --type=renderer", cfg))
        # Shell protection & runaway exemption
        self.assertTrue(w.is_protected(9999, "bash", "/bin/bash", {"protected_commands": ["bash"]}, rss_mb=50))
        self.assertFalse(w.is_protected(9999, "bash", "bash test_script.sh", {"protected_commands": ["bash"]}, rss_mb=2500))


class TestStabilityEvaluation(unittest.TestCase):
    def setUp(self):
        self.cfg = {
            "thresholds": {
                "load_warning": 20.0,
                "load_critical": 35.0,
                "load_emergency": 60.0,
                "ram_warning_pct": 80.0,
                "ram_critical_pct": 90.0,
                "swap_warning_pct": 75.0,
                "swap_critical_pct": 85.0,
                "process_rss_warning_mb": 2000,
                "process_rss_critical_mb": 3000,
            },
            "protected_commands": ["sshd", "tmux"]
        }

    def test_green_tier(self):
        metrics = {"load_1m": 2.5, "ram_used_pct": 30.0, "swap_used_pct": 10.0}
        procs = [
            {"pid": 101, "cmdline": "muse-bin", "rss_mb": 400, "cpu_pct": 5.0, "is_protected": False, "nice": 0}
        ]
        tier, reasons, actions = w.evaluate_stability(metrics, procs, self.cfg)
        self.assertEqual(tier, "GREEN")
        self.assertEqual(reasons, [])
        self.assertEqual(actions, [])

    def test_yellow_tier_elevated_load(self):
        metrics = {"load_1m": 22.5, "ram_used_pct": 50.0, "swap_used_pct": 20.0}
        procs = [
            {"pid": 102, "cmdline": "python worker", "rss_mb": 500, "cpu_pct": 90.0, "is_protected": False, "nice": 0}
        ]
        tier, reasons, actions = w.evaluate_stability(metrics, procs, self.cfg)
        self.assertEqual(tier, "YELLOW")
        self.assertTrue(any("Elevated load/memory" in r for r in reasons))
        self.assertEqual(len(actions), 1)
        self.assertEqual(actions[0]["action"], "renice")
        self.assertEqual(actions[0]["pid"], 102)

    def test_orange_tier_pause_leaky_process(self):
        metrics = {"load_1m": 2.0, "ram_used_pct": 40.0, "swap_used_pct": 10.0}
        procs = [
            {"pid": 202, "cmdline": "muse-bin leaky", "rss_mb": 3400, "cpu_pct": 10.0, "is_protected": False, "nice": 0}
        ]
        tier, reasons, actions = w.evaluate_stability(metrics, procs, self.cfg)
        self.assertEqual(tier, "ORANGE")
        self.assertTrue(any("exceeded critical RSS" in r for r in reasons))
        self.assertEqual(len(actions), 1)
        self.assertEqual(actions[0]["action"], "pause")
        self.assertEqual(actions[0]["pid"], 202)

    def test_red_emergency_tier(self):
        metrics = {"load_1m": 75.0, "ram_used_pct": 96.0, "swap_used_pct": 94.0}
        procs = [
            {"pid": 301, "cmdline": "muse-bin heavy", "rss_mb": 1800, "cpu_pct": 50.0, "is_protected": False, "nice": 0},
            {"pid": 302, "cmdline": "sshd daemon", "rss_mb": 2000, "cpu_pct": 2.0, "is_protected": True, "nice": 0}
        ]
        tier, reasons, actions = w.evaluate_stability(metrics, procs, self.cfg)
        self.assertEqual(tier, "RED")
        self.assertTrue(any("Emergency host pressure" in r for r in reasons))
        pause_pids = [a["pid"] for a in actions if a["action"] == "pause"]
        self.assertIn(301, pause_pids)
        self.assertNotIn(302, pause_pids)


class TestSocketIsolation(unittest.TestCase):
    def test_distinguishes_user_from_box_launched_agents(self):
        # PID 555 is an automated job on default socket
        # PID 666 is an agent on the correct fleet socket
        # PID 777 is a user-launched interactive agent on default socket
        procs = [
            {"pid": 555, "cmdline": "muse-bin auto-work sweep", "tmux_sock": "/tmp/tmux-1000/default", "is_protected": False},
            {"pid": 666, "cmdline": "muse-bin auto-work sweep", "tmux_sock": "/tmp/tmux-muse.sock", "is_protected": False},
            {"pid": 777, "cmdline": "muse-bin interactive chat", "tmux_sock": "/tmp/tmux-1000/default", "is_protected": False},
        ]
        box_sessions = {"auto-work": {}}
        violations, allowed = w.check_socket_isolation_violations(procs, box_sessions=box_sessions)
        self.assertEqual(len(violations), 1)
        self.assertEqual(violations[0]["pid"], 555)
        self.assertEqual(len(allowed), 1)
        self.assertEqual(allowed[0]["pid"], 777)


class TestPauseResumeAndExpiry(unittest.TestCase):
    @mock.patch("os.kill")
    def test_pause_and_state_persistence(self, mock_kill):
        with tempfile.TemporaryDirectory() as td:
            state_file = Path(td) / "paused.json"
            actions = [{"action": "pause", "pid": 4321, "cmd": "muse-bin", "reason": "RSS high"}]
            executed = w.execute_actions(actions, state_file=state_file, dry_run=False)
            self.assertEqual(len(executed), 1)
            mock_kill.assert_called_once_with(4321, 19)  # SIGSTOP = 19
            self.assertTrue(state_file.exists())
            data = json.loads(state_file.read_text())
            self.assertIn("4321", data)

    @mock.patch("os.kill")
    def test_reconcile_expired_pause(self, mock_kill):
        with tempfile.TemporaryDirectory() as td:
            state_file = Path(td) / "paused.json"
            now = w.now_epoch()
            state_data = {
                "4321": {"pid": 4321, "cmd": "muse-bin", "paused_at_epoch": now - 70}
            }
            state_file.write_text(json.dumps(state_data))

            expired = w.reconcile_paused_processes(state_file=state_file, dry_run=False)
            self.assertEqual(len(expired), 1)
            self.assertEqual(expired[0]["action"], "cull_expired")
            self.assertEqual(expired[0]["pid"], 4321)
            mock_kill.assert_called_once_with(4321, 15)  # SIGTERM = 15
            remaining = json.loads(state_file.read_text())
            self.assertNotIn("4321", remaining)

    @mock.patch("os.kill")
    def test_resume_process(self, mock_kill):
        with tempfile.TemporaryDirectory() as td:
            state_file = Path(td) / "paused.json"
            state_file.write_text(json.dumps({"4321": {"pid": 4321}}))
            res = w.resume_process(4321, state_file=state_file)
            self.assertTrue(res["success"])
            mock_kill.assert_called_once_with(4321, 18)  # SIGCONT = 18
            remaining = json.loads(state_file.read_text())
            self.assertNotIn("4321", remaining)


if __name__ == "__main__":
    unittest.main()