mirror of
https://github.com/Monadical-SAS/cubbi.git
synced 2025-12-20 12:19:07 +00:00
d750e64608
Add mcps field to Session model to track active MCP servers and populate it from container labels in ContainerManager. Enhance MCP remove command to warn when removing servers used by active sessions. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Claude <noreply@anthropic.com>
506 lines
16 KiB
Python
506 lines
16 KiB
Python
"""
|
|
Tests for the MCP server management commands.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import patch
|
|
from cubbi.cli import app
|
|
|
|
|
|
def test_mcp_list_empty(cli_runner, patched_config_manager):
|
|
"""Test the 'cubbi mcp list' command with no MCPs configured."""
|
|
# Make sure mcps is empty
|
|
patched_config_manager.set("mcps", [])
|
|
|
|
with patch("cubbi.cli.mcp_manager.list_mcps") as mock_list_mcps:
|
|
mock_list_mcps.return_value = []
|
|
|
|
result = cli_runner.invoke(app, ["mcp", "list"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "No MCP servers configured" in result.stdout
|
|
|
|
|
|
def test_mcp_add_remote(cli_runner, patched_config_manager):
|
|
"""Test adding a remote MCP server and listing it."""
|
|
# Add a remote MCP server
|
|
result = cli_runner.invoke(
|
|
app,
|
|
[
|
|
"mcp",
|
|
"add-remote",
|
|
"test-remote-mcp",
|
|
"http://mcp-server.example.com/sse",
|
|
"--header",
|
|
"Authorization=Bearer test-token",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert "Added remote MCP server" in result.stdout
|
|
|
|
# List MCP servers
|
|
result = cli_runner.invoke(app, ["mcp", "list"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "test-remote-mcp" in result.stdout
|
|
assert "remote" in result.stdout
|
|
# Check partial URL since it may be truncated in the table display
|
|
assert "http://mcp-se" in result.stdout # Truncated in table view
|
|
|
|
|
|
def test_mcp_add(cli_runner, patched_config_manager):
|
|
"""Test adding a proxy-based MCP server and listing it."""
|
|
# Add a Docker MCP server
|
|
result = cli_runner.invoke(
|
|
app,
|
|
[
|
|
"mcp",
|
|
"add",
|
|
"test-docker-mcp",
|
|
"mcp/github:latest",
|
|
"--command",
|
|
"github-mcp",
|
|
"--env",
|
|
"GITHUB_TOKEN=test-token",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert "Added MCP server" in result.stdout
|
|
|
|
# List MCP servers
|
|
result = cli_runner.invoke(app, ["mcp", "list"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "test-docker-mcp" in result.stdout
|
|
assert "proxy" in result.stdout # It's a proxy-based MCP
|
|
assert "mcp/github:la" in result.stdout # Truncated in table view
|
|
|
|
|
|
def test_mcp_remove(cli_runner, patched_config_manager):
|
|
"""Test removing an MCP server."""
|
|
# Add a remote MCP server
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-mcp",
|
|
"type": "remote",
|
|
"url": "http://test-server.com/sse",
|
|
"headers": {"Authorization": "Bearer test-token"},
|
|
}
|
|
],
|
|
)
|
|
|
|
# Mock the container_manager.list_sessions to return sessions without MCPs
|
|
with patch("cubbi.cli.container_manager.list_sessions") as mock_list_sessions:
|
|
mock_list_sessions.return_value = []
|
|
|
|
# Mock the remove_mcp method
|
|
with patch("cubbi.cli.mcp_manager.remove_mcp") as mock_remove_mcp:
|
|
# Make remove_mcp return True (successful removal)
|
|
mock_remove_mcp.return_value = True
|
|
|
|
# Remove the MCP server
|
|
result = cli_runner.invoke(app, ["mcp", "remove", "test-mcp"])
|
|
|
|
# Just check it ran successfully with exit code 0
|
|
assert result.exit_code == 0
|
|
assert "Removed MCP server 'test-mcp'" in result.stdout
|
|
|
|
|
|
def test_mcp_remove_with_active_sessions(cli_runner, patched_config_manager):
|
|
"""Test removing an MCP server that is used by active sessions."""
|
|
from cubbi.models import Session, SessionStatus
|
|
|
|
# Add a remote MCP server
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-mcp",
|
|
"type": "remote",
|
|
"url": "http://test-server.com/sse",
|
|
"headers": {"Authorization": "Bearer test-token"},
|
|
}
|
|
],
|
|
)
|
|
|
|
# Create mock sessions that use the MCP
|
|
mock_sessions = [
|
|
Session(
|
|
id="session-1",
|
|
name="test-session-1",
|
|
image="goose",
|
|
status=SessionStatus.RUNNING,
|
|
container_id="container-1",
|
|
mcps=["test-mcp", "other-mcp"],
|
|
),
|
|
Session(
|
|
id="session-2",
|
|
name="test-session-2",
|
|
image="goose",
|
|
status=SessionStatus.RUNNING,
|
|
container_id="container-2",
|
|
mcps=["other-mcp"], # This one doesn't use test-mcp
|
|
),
|
|
Session(
|
|
id="session-3",
|
|
name="test-session-3",
|
|
image="goose",
|
|
status=SessionStatus.RUNNING,
|
|
container_id="container-3",
|
|
mcps=["test-mcp"], # This one uses test-mcp
|
|
),
|
|
]
|
|
|
|
# Mock the container_manager.list_sessions to return our sessions
|
|
with patch("cubbi.cli.container_manager.list_sessions") as mock_list_sessions:
|
|
mock_list_sessions.return_value = mock_sessions
|
|
|
|
# Mock the remove_mcp method
|
|
with patch("cubbi.cli.mcp_manager.remove_mcp") as mock_remove_mcp:
|
|
# Make remove_mcp return True (successful removal)
|
|
mock_remove_mcp.return_value = True
|
|
|
|
# Remove the MCP server
|
|
result = cli_runner.invoke(app, ["mcp", "remove", "test-mcp"])
|
|
|
|
# Check it ran successfully with exit code 0
|
|
assert result.exit_code == 0
|
|
assert "Removed MCP server 'test-mcp'" in result.stdout
|
|
# Check warning about affected sessions
|
|
assert (
|
|
"Warning: Found 2 active sessions using MCP 'test-mcp'" in result.stdout
|
|
)
|
|
assert "session-1" in result.stdout
|
|
assert "session-3" in result.stdout
|
|
# session-2 should not be mentioned since it doesn't use test-mcp
|
|
assert "session-2" not in result.stdout
|
|
|
|
|
|
def test_mcp_remove_nonexistent(cli_runner, patched_config_manager):
|
|
"""Test removing a non-existent MCP server."""
|
|
# No MCPs configured
|
|
patched_config_manager.set("mcps", [])
|
|
|
|
# Mock the container_manager.list_sessions to return empty list
|
|
with patch("cubbi.cli.container_manager.list_sessions") as mock_list_sessions:
|
|
mock_list_sessions.return_value = []
|
|
|
|
# Mock the remove_mcp method to return False (MCP not found)
|
|
with patch("cubbi.cli.mcp_manager.remove_mcp") as mock_remove_mcp:
|
|
mock_remove_mcp.return_value = False
|
|
|
|
# Try to remove a non-existent MCP server
|
|
result = cli_runner.invoke(app, ["mcp", "remove", "nonexistent-mcp"])
|
|
|
|
# Check it ran successfully but reported not found
|
|
assert result.exit_code == 0
|
|
assert "MCP server 'nonexistent-mcp' not found" in result.stdout
|
|
|
|
|
|
def test_session_mcps_attribute():
|
|
"""Test that Session model has mcps attribute and can be populated correctly."""
|
|
from cubbi.models import Session, SessionStatus
|
|
|
|
# Test that Session can be created with mcps attribute
|
|
session = Session(
|
|
id="test-session",
|
|
name="test-session",
|
|
image="goose",
|
|
status=SessionStatus.RUNNING,
|
|
container_id="test-container",
|
|
mcps=["mcp1", "mcp2"],
|
|
)
|
|
|
|
assert session.mcps == ["mcp1", "mcp2"]
|
|
|
|
# Test that Session can be created with empty mcps list
|
|
session_empty = Session(
|
|
id="test-session-2",
|
|
name="test-session-2",
|
|
image="goose",
|
|
status=SessionStatus.RUNNING,
|
|
container_id="test-container-2",
|
|
)
|
|
|
|
assert session_empty.mcps == [] # Should default to empty list
|
|
|
|
|
|
def test_session_mcps_from_container_labels():
|
|
"""Test that Session mcps are correctly populated from container labels."""
|
|
from unittest.mock import Mock
|
|
from cubbi.container import ContainerManager
|
|
|
|
# Mock a container with MCP labels
|
|
mock_container = Mock()
|
|
mock_container.id = "test-container-id"
|
|
mock_container.status = "running"
|
|
mock_container.labels = {
|
|
"cubbi.session": "true",
|
|
"cubbi.session.id": "test-session",
|
|
"cubbi.session.name": "test-session-name",
|
|
"cubbi.image": "goose",
|
|
"cubbi.mcps": "mcp1,mcp2,mcp3", # Test with multiple MCPs
|
|
}
|
|
mock_container.attrs = {"NetworkSettings": {"Ports": {}}}
|
|
|
|
# Mock Docker client
|
|
mock_client = Mock()
|
|
mock_client.containers.list.return_value = [mock_container]
|
|
|
|
# Create container manager with mocked client
|
|
with patch("cubbi.container.docker.from_env") as mock_docker:
|
|
mock_docker.return_value = mock_client
|
|
mock_client.ping.return_value = True
|
|
|
|
container_manager = ContainerManager()
|
|
sessions = container_manager.list_sessions()
|
|
|
|
assert len(sessions) == 1
|
|
session = sessions[0]
|
|
assert session.id == "test-session"
|
|
assert session.mcps == ["mcp1", "mcp2", "mcp3"]
|
|
|
|
|
|
def test_session_mcps_from_empty_container_labels():
|
|
"""Test that Session mcps are correctly handled when container has no MCP labels."""
|
|
from unittest.mock import Mock
|
|
from cubbi.container import ContainerManager
|
|
|
|
# Mock a container without MCP labels
|
|
mock_container = Mock()
|
|
mock_container.id = "test-container-id"
|
|
mock_container.status = "running"
|
|
mock_container.labels = {
|
|
"cubbi.session": "true",
|
|
"cubbi.session.id": "test-session",
|
|
"cubbi.session.name": "test-session-name",
|
|
"cubbi.image": "goose",
|
|
# No cubbi.mcps label
|
|
}
|
|
mock_container.attrs = {"NetworkSettings": {"Ports": {}}}
|
|
|
|
# Mock Docker client
|
|
mock_client = Mock()
|
|
mock_client.containers.list.return_value = [mock_container]
|
|
|
|
# Create container manager with mocked client
|
|
with patch("cubbi.container.docker.from_env") as mock_docker:
|
|
mock_docker.return_value = mock_client
|
|
mock_client.ping.return_value = True
|
|
|
|
container_manager = ContainerManager()
|
|
sessions = container_manager.list_sessions()
|
|
|
|
assert len(sessions) == 1
|
|
session = sessions[0]
|
|
assert session.id == "test-session"
|
|
assert session.mcps == [] # Should be empty list when no MCPs
|
|
|
|
|
|
@pytest.mark.requires_docker
|
|
def test_mcp_status(cli_runner, patched_config_manager, mock_container_manager):
|
|
"""Test the MCP status command."""
|
|
# Add a Docker MCP
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
"env": {"TEST_ENV": "test-value"},
|
|
}
|
|
],
|
|
)
|
|
|
|
# First mock get_mcp to return our MCP config
|
|
with patch("cubbi.cli.mcp_manager.get_mcp") as mock_get_mcp:
|
|
mock_get_mcp.return_value = {
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
"env": {"TEST_ENV": "test-value"},
|
|
}
|
|
|
|
# Then mock the get_mcp_status method
|
|
with patch("cubbi.cli.mcp_manager.get_mcp_status") as mock_get_status:
|
|
mock_get_status.return_value = {
|
|
"status": "running",
|
|
"container_id": "test-container-id",
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"ports": {"8080/tcp": 8080},
|
|
"created": "2023-01-01T00:00:00Z",
|
|
}
|
|
|
|
# Check MCP status
|
|
result = cli_runner.invoke(app, ["mcp", "status", "test-docker-mcp"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "test-docker-mcp" in result.stdout
|
|
assert "running" in result.stdout
|
|
assert "mcp/test:latest" in result.stdout
|
|
|
|
|
|
@pytest.mark.requires_docker
|
|
def test_mcp_start(cli_runner, patched_config_manager, mock_container_manager):
|
|
"""Test starting an MCP server."""
|
|
# Add a Docker MCP
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
}
|
|
],
|
|
)
|
|
|
|
# Mock the start operation
|
|
mock_container_manager.start_mcp.return_value = {
|
|
"container_id": "test-container-id",
|
|
"status": "running",
|
|
}
|
|
|
|
# Start the MCP
|
|
result = cli_runner.invoke(app, ["mcp", "start", "test-docker-mcp"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Started MCP server" in result.stdout
|
|
assert "test-docker-mcp" in result.stdout
|
|
|
|
|
|
@pytest.mark.requires_docker
|
|
def test_mcp_stop(cli_runner, patched_config_manager, mock_container_manager):
|
|
"""Test stopping an MCP server."""
|
|
# Add a Docker MCP
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
}
|
|
],
|
|
)
|
|
|
|
# Mock the stop operation
|
|
mock_container_manager.stop_mcp.return_value = True
|
|
|
|
# Stop the MCP
|
|
result = cli_runner.invoke(app, ["mcp", "stop", "test-docker-mcp"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Stopped and removed MCP server" in result.stdout
|
|
assert "test-docker-mcp" in result.stdout
|
|
|
|
|
|
@pytest.mark.requires_docker
|
|
def test_mcp_restart(cli_runner, patched_config_manager, mock_container_manager):
|
|
"""Test restarting an MCP server."""
|
|
# Add a Docker MCP
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
}
|
|
],
|
|
)
|
|
|
|
# Mock the restart operation
|
|
mock_container_manager.restart_mcp.return_value = {
|
|
"container_id": "test-container-id",
|
|
"status": "running",
|
|
}
|
|
|
|
# Restart the MCP
|
|
result = cli_runner.invoke(app, ["mcp", "restart", "test-docker-mcp"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Restarted MCP server" in result.stdout
|
|
assert "test-docker-mcp" in result.stdout
|
|
|
|
|
|
@pytest.mark.requires_docker
|
|
def test_mcp_logs(cli_runner, patched_config_manager, mock_container_manager):
|
|
"""Test viewing MCP server logs."""
|
|
# Add a Docker MCP
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-docker-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
}
|
|
],
|
|
)
|
|
|
|
# Mock the logs operation
|
|
with patch("cubbi.cli.mcp_manager.get_mcp_logs") as mock_get_logs:
|
|
mock_get_logs.return_value = "Test log output"
|
|
|
|
# View MCP logs
|
|
result = cli_runner.invoke(app, ["mcp", "logs", "test-docker-mcp"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Test log output" in result.stdout
|
|
|
|
|
|
def test_session_with_mcp(cli_runner, patched_config_manager, mock_container_manager):
|
|
"""Test creating a session with an MCP server attached."""
|
|
# Add an MCP server
|
|
patched_config_manager.set(
|
|
"mcps",
|
|
[
|
|
{
|
|
"name": "test-mcp",
|
|
"type": "docker",
|
|
"image": "mcp/test:latest",
|
|
"command": "test-command",
|
|
}
|
|
],
|
|
)
|
|
|
|
# Mock the session creation with MCP
|
|
from cubbi.models import Session, SessionStatus
|
|
|
|
# timestamp no longer needed since we don't use created_at in Session
|
|
mock_container_manager.create_session.return_value = Session(
|
|
id="test-session-id",
|
|
name="test-session",
|
|
image="goose",
|
|
status=SessionStatus.RUNNING,
|
|
container_id="test-container-id",
|
|
ports={},
|
|
)
|
|
|
|
# Create a session with MCP
|
|
result = cli_runner.invoke(app, ["session", "create", "--mcp", "test-mcp"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Session created successfully" in result.stdout
|
|
assert "test-session" in result.stdout
|
|
# Check that the create_session was called with the mcp parameter
|
|
assert mock_container_manager.create_session.called
|
|
# The keyword arguments are in the second element of call_args
|
|
kwargs = mock_container_manager.create_session.call_args[1]
|
|
assert "mcp" in kwargs
|
|
assert "test-mcp" in kwargs["mcp"]
|