304 lines
11 KiB
Python
304 lines
11 KiB
Python
"""Tests for daemon health check and reconnection logic."""
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from beads_mcp.bd_client import BdError
|
|
from beads_mcp.bd_daemon_client import (
|
|
BdDaemonClient,
|
|
DaemonConnectionError,
|
|
DaemonError,
|
|
DaemonNotRunningError,
|
|
)
|
|
from beads_mcp.tools import _get_client, _health_check_client, _reconnect_client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_daemon_client_ping_success():
|
|
"""Test successful ping to daemon."""
|
|
client = BdDaemonClient(socket_path="/tmp/bd.sock", working_dir="/tmp/test")
|
|
|
|
with patch.object(client, '_send_request', new_callable=AsyncMock) as mock_send:
|
|
mock_send.return_value = {"message": "pong", "version": "0.9.10"}
|
|
|
|
result = await client.ping()
|
|
|
|
assert result["message"] == "pong"
|
|
assert result["version"] == "0.9.10"
|
|
mock_send.assert_called_once_with("ping", {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_daemon_client_ping_connection_error():
|
|
"""Test ping when daemon connection fails."""
|
|
client = BdDaemonClient(socket_path="/tmp/bd.sock", working_dir="/tmp/test")
|
|
|
|
with patch.object(client, '_send_request', new_callable=AsyncMock) as mock_send:
|
|
mock_send.side_effect = DaemonConnectionError("Connection failed")
|
|
|
|
with pytest.raises(DaemonConnectionError):
|
|
await client.ping()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_daemon_client_health_success():
|
|
"""Test successful health check to daemon."""
|
|
client = BdDaemonClient(socket_path="/tmp/bd.sock", working_dir="/tmp/test")
|
|
|
|
with patch.object(client, '_send_request', new_callable=AsyncMock) as mock_send:
|
|
mock_send.return_value = {
|
|
"status": "healthy",
|
|
"version": "0.9.10",
|
|
"uptime": 123.45,
|
|
"db_response_time_ms": 2.5,
|
|
"active_connections": 3,
|
|
"memory_bytes": 104857600,
|
|
}
|
|
|
|
result = await client.health()
|
|
|
|
assert result["status"] == "healthy"
|
|
assert result["version"] == "0.9.10"
|
|
assert result["uptime"] == 123.45
|
|
mock_send.assert_called_once_with("health", {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_daemon_client_health_unhealthy():
|
|
"""Test health check when daemon is unhealthy."""
|
|
client = BdDaemonClient(socket_path="/tmp/bd.sock", working_dir="/tmp/test")
|
|
|
|
with patch.object(client, '_send_request', new_callable=AsyncMock) as mock_send:
|
|
mock_send.return_value = {
|
|
"status": "unhealthy",
|
|
"error": "Database connection failed",
|
|
}
|
|
|
|
result = await client.health()
|
|
|
|
assert result["status"] == "unhealthy"
|
|
assert "error" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_client_daemon_client_healthy():
|
|
"""Test health check for healthy daemon client."""
|
|
client = BdDaemonClient(socket_path="/tmp/bd.sock", working_dir="/tmp/test")
|
|
|
|
with patch.object(client, 'ping', new_callable=AsyncMock) as mock_ping:
|
|
mock_ping.return_value = {"message": "pong", "version": "0.9.10"}
|
|
|
|
result = await _health_check_client(client)
|
|
|
|
assert result is True
|
|
mock_ping.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_client_daemon_client_unhealthy():
|
|
"""Test health check for unhealthy daemon client."""
|
|
client = BdDaemonClient(socket_path="/tmp/bd.sock", working_dir="/tmp/test")
|
|
|
|
with patch.object(client, 'ping', new_callable=AsyncMock) as mock_ping:
|
|
mock_ping.side_effect = DaemonConnectionError("Connection failed")
|
|
|
|
result = await _health_check_client(client)
|
|
|
|
assert result is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_client_cli_client():
|
|
"""Test health check for CLI client (always returns True)."""
|
|
from beads_mcp.bd_client import BdClient
|
|
|
|
client = BdClient(bd_path="/usr/bin/bd", beads_db="/tmp/test.db")
|
|
|
|
result = await _health_check_client(client)
|
|
|
|
# CLI clients don't have ping, so they're always considered healthy
|
|
assert result is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reconnect_client_success():
|
|
"""Test successful reconnection after failure."""
|
|
from beads_mcp.bd_client import create_bd_client
|
|
|
|
with (
|
|
patch('beads_mcp.tools.create_bd_client') as mock_create,
|
|
patch('beads_mcp.tools._health_check_client', new_callable=AsyncMock) as mock_health,
|
|
patch('beads_mcp.tools._register_client_for_cleanup') as mock_register,
|
|
):
|
|
mock_client = MagicMock()
|
|
mock_create.return_value = mock_client
|
|
mock_health.return_value = True
|
|
|
|
result = await _reconnect_client("/tmp/test")
|
|
|
|
assert result == mock_client
|
|
mock_create.assert_called_once_with(prefer_daemon=True, working_dir="/tmp/test")
|
|
mock_register.assert_called_once_with(mock_client)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reconnect_client_retry_with_backoff():
|
|
"""Test reconnection with exponential backoff on failure."""
|
|
# Need to patch asyncio.sleep in the actual module where it's called
|
|
import beads_mcp.tools as tools_module
|
|
|
|
with (
|
|
patch.object(tools_module, 'create_bd_client') as mock_create,
|
|
patch.object(tools_module, '_health_check_client', new_callable=AsyncMock) as mock_health,
|
|
patch.object(tools_module, '_register_client_for_cleanup') as mock_register,
|
|
):
|
|
mock_client = MagicMock()
|
|
|
|
# Raise exception first two times, succeed third time
|
|
mock_create.side_effect = [
|
|
Exception("Connection failed"),
|
|
Exception("Connection failed"),
|
|
mock_client,
|
|
]
|
|
mock_health.return_value = True
|
|
|
|
# Mock asyncio.sleep to track calls
|
|
sleep_calls = []
|
|
async def mock_sleep(duration):
|
|
sleep_calls.append(duration)
|
|
# Don't actually sleep in tests
|
|
return
|
|
|
|
with patch.object(asyncio, 'sleep', side_effect=mock_sleep):
|
|
result = await _reconnect_client("/tmp/test", max_retries=3)
|
|
|
|
assert result == mock_client
|
|
assert mock_create.call_count == 3
|
|
assert len(sleep_calls) == 2
|
|
|
|
# Verify exponential backoff: 0.1s, 0.2s
|
|
assert sleep_calls[0] == 0.1
|
|
assert sleep_calls[1] == 0.2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reconnect_client_max_retries_exceeded():
|
|
"""Test reconnection failure after max retries."""
|
|
with (
|
|
patch('beads_mcp.tools.create_bd_client') as mock_create,
|
|
patch('beads_mcp.tools._health_check_client', new_callable=AsyncMock) as mock_health,
|
|
patch('asyncio.sleep', new_callable=AsyncMock),
|
|
):
|
|
mock_client = MagicMock()
|
|
mock_create.return_value = mock_client
|
|
mock_health.return_value = False # Always fail health check
|
|
|
|
with pytest.raises(BdError, match="Failed to connect to daemon after 3 attempts"):
|
|
await _reconnect_client("/tmp/test", max_retries=3)
|
|
|
|
assert mock_create.call_count == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_client_uses_cached_healthy_client(monkeypatch):
|
|
"""Test that _get_client returns cached client if healthy."""
|
|
from beads_mcp import tools
|
|
|
|
# Set up environment
|
|
monkeypatch.setenv("BEADS_WORKING_DIR", "/tmp/test")
|
|
|
|
mock_client = MagicMock()
|
|
mock_client._check_version = AsyncMock()
|
|
|
|
with (
|
|
patch('beads_mcp.tools._canonicalize_path', return_value="/tmp/test"),
|
|
patch('beads_mcp.tools._health_check_client', new_callable=AsyncMock) as mock_health,
|
|
):
|
|
mock_health.return_value = True
|
|
|
|
# Add mock client to pool and mark as version checked
|
|
tools._connection_pool["/tmp/test"] = mock_client
|
|
tools._version_checked.add("/tmp/test")
|
|
|
|
result = await _get_client()
|
|
|
|
assert result == mock_client
|
|
mock_health.assert_called_once_with(mock_client)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_client_reconnects_on_stale_connection(monkeypatch):
|
|
"""Test that _get_client reconnects when cached client is stale."""
|
|
from beads_mcp import tools
|
|
|
|
# Set up environment
|
|
monkeypatch.setenv("BEADS_WORKING_DIR", "/tmp/test")
|
|
|
|
old_client = MagicMock()
|
|
new_client = MagicMock()
|
|
new_client._check_version = AsyncMock()
|
|
|
|
with (
|
|
patch('beads_mcp.tools._canonicalize_path', return_value="/tmp/test"),
|
|
patch('beads_mcp.tools._health_check_client', new_callable=AsyncMock) as mock_health,
|
|
patch('beads_mcp.tools._reconnect_client', new_callable=AsyncMock) as mock_reconnect,
|
|
):
|
|
# First health check fails (stale), reconnect returns new client
|
|
mock_health.return_value = False
|
|
mock_reconnect.return_value = new_client
|
|
|
|
# Add old client to pool
|
|
tools._connection_pool["/tmp/test"] = old_client
|
|
tools._version_checked.add("/tmp/test")
|
|
|
|
result = await _get_client()
|
|
|
|
assert result == new_client
|
|
assert tools._connection_pool["/tmp/test"] == new_client
|
|
# Version check is performed after reconnect, so it's back in the set
|
|
assert "/tmp/test" in tools._version_checked
|
|
mock_reconnect.assert_called_once_with("/tmp/test")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_client_creates_new_client_if_not_cached(monkeypatch):
|
|
"""Test that _get_client creates new client if not in pool."""
|
|
from beads_mcp import tools
|
|
|
|
# Clear pool
|
|
tools._connection_pool.clear()
|
|
tools._version_checked.clear()
|
|
|
|
# Set up environment
|
|
monkeypatch.setenv("BEADS_WORKING_DIR", "/tmp/test")
|
|
|
|
mock_client = MagicMock()
|
|
mock_client._check_version = AsyncMock()
|
|
|
|
with (
|
|
patch('beads_mcp.tools._canonicalize_path', return_value="/tmp/test"),
|
|
patch('beads_mcp.tools.create_bd_client', return_value=mock_client) as mock_create,
|
|
patch('beads_mcp.tools._register_client_for_cleanup') as mock_register,
|
|
):
|
|
result = await _get_client()
|
|
|
|
assert result == mock_client
|
|
assert tools._connection_pool["/tmp/test"] == mock_client
|
|
mock_create.assert_called_once_with(prefer_daemon=True, working_dir="/tmp/test")
|
|
mock_register.assert_called_once_with(mock_client)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_client_no_workspace_error():
|
|
"""Test that _get_client raises error if no workspace is set."""
|
|
from beads_mcp import tools
|
|
|
|
# Clear context
|
|
tools.current_workspace.set(None)
|
|
|
|
with patch.dict('os.environ', {}, clear=True):
|
|
with pytest.raises(BdError, match="No workspace set"):
|
|
await _get_client()
|