box/tests/test_notifications.py

import unittest
import time
from unittest.mock import MagicMock
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "bin"))

import muse_tui

class TestNotificationCounters(unittest.TestCase):
    def setUp(self):
        self.dm = muse_tui.FleetDataManager()
        # Mock background loops from interfering with testing
        self.dm.running = False
        self.dm.read_msg_counts = {}
        # Isolate threads and history caches for tests
        self.dm.nodes = ["pip", "646"]
        self.dm.threads_cache = {"pip": [], "646": []}
        self.dm.history_cache = {}

    def test_unread_count_calculation(self):
        node = "pip"
        tid = "test-thread-123"
        key = (node, tid)

        # Initially no messages -> unread is 0
        self.assertEqual(self.dm.get_unread_count(node, tid), 0)

        # 5 messages arrive in cache
        self.dm.history_cache[key] = [{"text": f"msg {i}"} for i in range(5)]
        self.assertEqual(self.dm.get_unread_count(node, tid), 5)

        # Read 3 messages
        self.dm.mark_thread_read(node, tid, count=3)
        self.assertEqual(self.dm.get_unread_count(node, tid), 2)

        # Read all messages
        self.dm.mark_thread_read(node, tid)
        self.assertEqual(self.dm.get_unread_count(node, tid), 0)

    def test_node_and_fleet_unread_total(self):
        # Configure test threads for pip and 646
        self.dm.threads_cache["pip"] = [
            {"session_id": "pip-chat-1", "is_main": True},
            {"session_id": "pip-chat-2", "is_main": False},
        ]
        self.dm.threads_cache["646"] = [
            {"session_id": "646-chat-1", "is_main": True},
        ]

        # Populate message histories
        self.dm.history_cache[("pip", "pip-chat-1")] = [{"text": "1"}, {"text": "2"}]
        self.dm.history_cache[("pip", "pip-chat-2")] = [{"text": "a"}, {"text": "b"}, {"text": "c"}]
        self.dm.history_cache[("646", "646-chat-1")] = [{"text": "x"}]

        # Before reading: pip has 5, 646 has 1, total fleet is 6
        self.assertEqual(self.dm.get_node_unread_total("pip"), 5)
        self.assertEqual(self.dm.get_node_unread_total("646"), 1)
        self.assertEqual(self.dm.get_fleet_unread_total(), 6)

        # Read pip-chat-1
        self.dm.mark_thread_read("pip", "pip-chat-1")
        self.assertEqual(self.dm.get_node_unread_total("pip"), 3)
        self.assertEqual(self.dm.get_fleet_unread_total(), 4)

    def test_dwell_clearing_logic(self):
        tui = muse_tui.MuseTUI(MagicMock())
        tui.data.running = False
        active_node = "muse"
        active_tid = "muse-thread-abc"

        tui.data.active_node = active_node
        tui.data.active_thread_id = active_tid
        tui.data.history_cache[(active_node, active_tid)] = [{"text": "msg"}] * 4
        self.assertEqual(tui.data.get_unread_count(active_node, active_tid), 4)

        # Simulated dwell < 5.0 seconds
        tui.cur_chat_enter_time = time.time() - 2.0
        tui.last_dwell_check_target = (active_node, active_tid)
        # Should not clear yet
        if time.time() - tui.cur_chat_enter_time >= 5.0:
            tui.data.mark_thread_read(active_node, active_tid)
        self.assertEqual(tui.data.get_unread_count(active_node, active_tid), 4)

        # Simulated dwell >= 5.0 seconds
        tui.cur_chat_enter_time = time.time() - 5.5
        if time.time() - tui.cur_chat_enter_time >= 5.0:
            tui.data.mark_thread_read(active_node, active_tid)
        self.assertEqual(tui.data.get_unread_count(active_node, active_tid), 0)

    def test_agent_sorting_and_initial_selection(self):
        # Configure nodes
        self.dm.nodes = ["pip", "muse", "opm", "646"]
        self.dm.threads_cache = {
            "pip": [{"session_id": "p1"}],
            "muse": [{"session_id": "m1"}],
            "opm": [{"session_id": "o1"}],
            "646": [{"session_id": "c1"}],
        }
        # Give opm 5 unread messages, muse 2 unreads, pip 0, 646 0
        self.dm.history_cache = {
            ("opm", "o1"): [{"text": f"msg {i}"} for i in range(5)],
            ("muse", "m1"): [{"text": f"msg {i}"} for i in range(2)],
        }
        self.dm.agent_interactions = {"pip": 100.0, "646": 200.0, "muse": 50.0, "opm": 0.0}

        self.dm.sort_nodes()
        # opm should be first (highest unread = 5)
        # muse should be second (unread = 2)
        # 646 should be third (interaction = 200.0)
        # pip should be fourth (interaction = 100.0)
        self.assertEqual(self.dm.nodes, ["opm", "muse", "646", "pip"])

        # Default initial_node=None in MuseTUI picks nodes[0] (top of list) and its Main Chat
        tui = muse_tui.MuseTUI(MagicMock(), initial_node=None)
        tui.data.running = False
        self.assertEqual(tui.data.active_node, tui.data.nodes[0])
        self.assertIsNotNone(tui.data.active_thread_id)
        top_node = tui.data.nodes[0]
        top_threads = {t.get("session_id") for t in tui.data.threads_cache.get(top_node, [])}
        self.assertIn(tui.data.active_thread_id, top_threads)

    def test_reply_workflow(self):
        tui = muse_tui.MuseTUI(MagicMock())
        tui.data.running = False

        msg = {
            "role": "assistant",
            "message_id": "msg-987654321",
            "seq": 42,
            "text": "Please confirm deployment parameters before we proceed.",
        }

        # 1. Trigger reply
        ok = tui.trigger_reply_to_message(msg)
        self.assertTrue(ok)
        self.assertIsNotNone(tui.reply_target)
        self.assertEqual(tui.reply_target["seq"], 42)
        self.assertEqual(tui.reply_target["role"], "assistant")
        self.assertEqual(tui.editor_mode, "INSERT")

        # 2. Cancel reply with Esc sequence
        tui._read_escape_sequence = MagicMock(return_value=[])  # standalone Esc
        tui._handle_escape_sequence()
        self.assertIsNone(tui.reply_target)
        self.assertEqual(tui.editor_mode, "NORMAL")

        # 3. Trigger reply again and execute input line
        tui.trigger_reply_to_message(msg)
        dispatched_messages = []
        tui._async_send_message = lambda node, tid, text: dispatched_messages.append((node, tid, text))

        tui._execute_input_line("Confirmed, deploy now.")
        self.assertIsNone(tui.reply_target)
        self.assertEqual(len(dispatched_messages), 1)
        node, tid, sent_text = dispatched_messages[0]
        self.assertIn("> Replying to ASSISTANT (seq:42):", sent_text)
        self.assertIn("Please confirm deployment parameters before we proceed.", sent_text)
        self.assertIn("Confirmed, deploy now.", sent_text)

    def test_sidechat_notifications_and_transcript_divider(self):
        tui = muse_tui.MuseTUI(MagicMock())
        tui.data.running = False
        active_node = "muse"
        active_tid = "sidechat-worker-42"
        tui.data.active_node = active_node
        tui.data.active_thread_id = active_tid

        tui.data.threads_cache[active_node] = [
            {"session_id": "main-chat", "is_main": True, "title": "Main Chat"},
            {"session_id": active_tid, "is_main": False, "title": "Worker Task 42"},
            {"session_id": "sidechat-worker-43", "is_main": False, "title": "Worker Task 43"},
        ]

        # 3 total messages, read count is 1 -> 2 unread messages in active sidechat
        tui.data.history_cache[(active_node, active_tid)] = [
            {"role": "user", "text": "initial prompt", "message_id": "m1"},
            {"role": "assistant", "text": "update 1", "message_id": "m2"},
            {"role": "assistant", "text": "update 2", "message_id": "m3"},
        ]
        tui.data.read_msg_counts[(active_node, active_tid)] = 1

        self.assertEqual(tui.data.get_unread_count(active_node, active_tid), 2)
        self.assertEqual(tui.data.get_node_unread_total(active_node), 2)

        # Mock safe_addstr to capture rendered strings
        rendered_texts = []
        tui.safe_addstr = lambda win, y, x, text, attr=0: rendered_texts.append(text)

        # 1. Render sidebar
        tui._render_sidebar(0, 0, 30, 32)
        # Should contain ●2 inline notification badge
        has_badge = any("●2" in s for s in rendered_texts)
        self.assertTrue(has_badge, "Sidechats list row or header must render inline ●2 unread badge")

        # 2. Render transcript
        rendered_texts.clear()
        tui._render_transcript(0, 32, 30, 80)
        # Should render the unread divider banner or unread pill
        has_div = any("2 NEW UNREAD MESSAGE" in s for s in rendered_texts)
        has_hdr_pill = any("[● 2 new]" in s for s in rendered_texts)
        self.assertTrue(has_div or has_hdr_pill, "Transcript must display unread banner or pill")

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