Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 19 additions & 21 deletions tests/_mcp/server/prompts/test_notebooks_prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from unittest.mock import Mock

from marimo._mcp.server._prompts.prompts.notebooks import ActiveNotebooks
from marimo._server.sessions import Session


def test_active_notebooks_metadata():
Expand All @@ -21,7 +20,7 @@ def test_active_notebooks_metadata():
def test_active_notebooks_no_sessions():
"""Test output when no sessions are active."""
context = Mock()
context.session_manager.sessions = {}
context.get_active_sessions_internal.return_value = []

prompt = ActiveNotebooks(context=context)
messages = prompt.handle()
Expand All @@ -38,41 +37,40 @@ def test_active_notebooks_with_sessions():
"""Test output with active sessions."""
context = Mock()

# Mock session with file path
session1 = Mock(spec=Session)
session1.app_file_manager = Mock()
session1.app_file_manager.filename = "/path/to/notebook.py"
# Mock active session objects
active_session1 = Mock()
active_session1.session_id = "session_1"
active_session1.path = "/path/to/notebook.py"

# Mock session without file path
session2 = Mock(spec=Session)
session2.app_file_manager = Mock()
session2.app_file_manager.filename = None
active_session2 = Mock()
active_session2.session_id = "session_2"
active_session2.path = None

context.session_manager.sessions = {
"session_1": session1,
"session_2": session2,
}
context.get_active_sessions_internal.return_value = [
active_session1,
active_session2,
]

prompt = ActiveNotebooks(context=context)
messages = prompt.handle()

assert len(messages) == 2
assert len(messages) == 3 # 2 sessions + 1 action message

# Check first message (with file path)
assert messages[0].role == "user"
assert messages[0].content.type == "text"
assert "session_1" in messages[0].content.text
assert "/path/to/notebook.py" in messages[0].content.text
assert (
"Use this session_id when calling MCP tools"
in messages[0].content.text
)

# Check second message (without file path)
assert messages[1].role == "user"
assert messages[1].content.type == "text"
assert "session_2" in messages[1].content.text

# Check action message
assert messages[2].role == "user"
assert messages[2].content.type == "text"
assert (
"Use this session_id when calling MCP tools"
in messages[1].content.text
"Use these session_ids when calling marimo MCP tools"
in messages[2].content.text
)
Loading