dbt-mcp

"""Unit tests for LocalLSPConnectionProvider class."""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from dbt_mcp.lsp.lsp_binary_manager import LspBinaryInfo
from dbt_mcp.lsp.providers.local_lsp_connection_provider import (
    LocalLSPConnectionProvider,
)
from dbt_mcp.lsp.lsp_connection import SocketLSPConnection


@pytest.fixture
def lsp_binary_info(tmp_path) -> LspBinaryInfo:
    """Create a test LSP binary info."""
    binary_path = tmp_path / "dbt-lsp"
    binary_path.touch()
    return LspBinaryInfo(path=str(binary_path), version="1.0.0")


@pytest.fixture
def project_dir(tmp_path) -> str:
    """Create a test project directory."""
    project_dir = tmp_path / "project"
    project_dir.mkdir()
    return str(project_dir)


class TestLocalLSPConnectionProvider:
    """Test LocalLSPConnectionProvider class."""

    @pytest.mark.asyncio
    async def test_get_connection_creates_connection_on_first_call(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that get_connection creates a new connection on first call."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Mock the SocketLSPConnection
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock()
        mock_connection.initialize = AsyncMock()

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ) as mock_conn_class:
            connection = await provider.get_connection()

            # Verify connection was created with correct arguments
            mock_conn_class.assert_called_once_with(
                binary_path=lsp_binary_info.path,
                args=[],
                cwd=project_dir,
            )

            # Verify connection lifecycle methods were called
            mock_connection.start.assert_called_once()
            mock_connection.initialize.assert_called_once()

            # Verify the connection is returned
            assert connection == mock_connection
            assert provider.lsp_connection == mock_connection

    @pytest.mark.asyncio
    async def test_get_connection_returns_existing_connection(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that get_connection returns the same connection on subsequent calls."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Mock the connection
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock()
        mock_connection.initialize = AsyncMock()

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ) as mock_conn_class:
            # First call
            connection1 = await provider.get_connection()

            # Second call
            connection2 = await provider.get_connection()

            # Third call
            connection3 = await provider.get_connection()

            # Connection should only be created once
            mock_conn_class.assert_called_once()
            mock_connection.start.assert_called_once()
            mock_connection.initialize.assert_called_once()

            # All calls should return the same instance
            assert connection1 is connection2
            assert connection2 is connection3
            assert connection1 is mock_connection

    @pytest.mark.asyncio
    async def test_get_connection_handles_start_failure(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that get_connection handles connection start failure."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Mock connection that fails to start
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock(side_effect=RuntimeError("Start failed"))
        mock_connection.initialize = AsyncMock()

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ):
            with pytest.raises(
                RuntimeError, match="Failed to establish LSP connection"
            ):
                await provider.get_connection()

            # Verify connection was cleaned up
            assert provider.lsp_connection is None

    @pytest.mark.asyncio
    async def test_get_connection_handles_initialize_failure(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that get_connection handles connection initialize failure."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Mock connection that fails to initialize
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock()
        mock_connection.initialize = AsyncMock(
            side_effect=RuntimeError("Initialize failed")
        )

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ):
            with pytest.raises(
                RuntimeError, match="Failed to establish LSP connection"
            ):
                await provider.get_connection()

            # Verify connection was cleaned up
            assert provider.lsp_connection is None

    @pytest.mark.asyncio
    async def test_cleanup_connection_stops_connection(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that cleanup_connection properly stops the connection."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Setup a connection
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock()
        mock_connection.initialize = AsyncMock()
        mock_connection.stop = AsyncMock()

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ):
            # Create connection
            await provider.get_connection()
            assert provider.lsp_connection is not None

            # Cleanup
            await provider.cleanup_connection()

            # Verify stop was called
            mock_connection.stop.assert_called_once()

            # Verify connection was set to None
            assert provider.lsp_connection is None

    @pytest.mark.asyncio
    async def test_cleanup_connection_handles_no_connection(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that cleanup_connection handles no existing connection gracefully."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # No connection has been created
        assert provider.lsp_connection is None

        # Should not raise
        await provider.cleanup_connection()

        assert provider.lsp_connection is None

    @pytest.mark.asyncio
    async def test_cleanup_connection_handles_stop_failure(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test that cleanup_connection handles stop failure gracefully."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Setup a connection that fails to stop
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock()
        mock_connection.initialize = AsyncMock()
        mock_connection.stop = AsyncMock(side_effect=RuntimeError("Stop failed"))

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ):
            # Create connection
            await provider.get_connection()

            # Cleanup should not raise
            await provider.cleanup_connection()

            # Verify stop was attempted
            mock_connection.stop.assert_called_once()

            # Connection should still be set to None
            assert provider.lsp_connection is None

    @pytest.mark.asyncio
    async def test_connection_lifecycle_integration(
        self, lsp_binary_info: LspBinaryInfo, project_dir: str
    ) -> None:
        """Test complete connection lifecycle (create, use, cleanup)."""
        provider = LocalLSPConnectionProvider(lsp_binary_info, project_dir)

        # Mock connection
        mock_connection = MagicMock(spec=SocketLSPConnection)
        mock_connection.start = AsyncMock()
        mock_connection.initialize = AsyncMock()
        mock_connection.stop = AsyncMock()

        with patch(
            "dbt_mcp.lsp.providers.local_lsp_connection_provider.SocketLSPConnection",
            return_value=mock_connection,
        ):
            # Create connection
            connection1 = await provider.get_connection()
            assert connection1 is mock_connection
            assert provider.lsp_connection is mock_connection

            # Get connection again (should return same)
            connection2 = await provider.get_connection()
            assert connection2 is connection1

            # Cleanup
            await provider.cleanup_connection()
            assert provider.lsp_connection is None

            # Get connection again (should create new)
            connection3 = await provider.get_connection()
            assert connection3 is mock_connection

            # Verify lifecycle methods were called correctly
            assert mock_connection.start.call_count == 2
            assert mock_connection.initialize.call_count == 2
            assert mock_connection.stop.call_count == 1