Fix langchain integration test failures + add DBSQL/streamable MCP tests (#369)

Author: dhruv0811

databricks_mcp/tests/integration_tests/conftest.py

@@ -180,6 +180,32 @@ def cached_vs_call_result(vs_mcp_client, cached_vs_tools_list):
     return vs_mcp_client.call_tool(tool.name, {param_name: "test"})
 
 
+# =============================================================================
+# DBSQL Fixtures
+# =============================================================================
+
+
+@pytest.fixture(scope="session")
+def dbsql_mcp_url(workspace_client):
+    """Construct MCP URL for the DBSQL server."""
+    base_url = workspace_client.config.host
+    return f"{base_url}/api/2.0/mcp/sql"
+
+
+@pytest.fixture(scope="session")
+def dbsql_mcp_client(dbsql_mcp_url, workspace_client):
+    """DatabricksMCPClient pointed at the DBSQL server."""
+    return DatabricksMCPClient(dbsql_mcp_url, workspace_client)
+
+
+@pytest.fixture(scope="session")
+def cached_dbsql_tools_list(dbsql_mcp_client):
+    """Cache the DBSQL list_tools() result."""
+    tools = dbsql_mcp_client.list_tools()
+    assert tools, "DBSQL list_tools() returned no tools"
+    return tools
+
+
 # =============================================================================
 # Genie Fixtures
 # =============================================================================

databricks_mcp/tests/integration_tests/test_mcp_core.py

@@ -8,12 +8,38 @@
 from __future__ import annotations
 
 import os
+from contextlib import asynccontextmanager
 
 import pytest
 from conftest import _skip_if_not_found
 from mcp.shared.exceptions import McpError
 from mcp.types import CallToolResult
 
+
+@asynccontextmanager
+async def raw_mcp_session(url, workspace_client):
+    """Create a raw MCP ClientSession using streamable_http_client with Databricks OAuth."""
+    import httpx
+    from mcp import ClientSession
+    from mcp.client.streamable_http import streamable_http_client
+
+    from databricks_mcp import DatabricksOAuthClientProvider
+
+    async with httpx.AsyncClient(
+        auth=DatabricksOAuthClientProvider(workspace_client),
+        follow_redirects=True,
+        timeout=httpx.Timeout(120.0, read=120.0),
+    ) as http_client:
+        async with streamable_http_client(url, http_client=http_client) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                yield session
+
+
 from databricks_mcp import DatabricksMCPClient
 
 pytestmark = pytest.mark.skipif(
@@ -126,6 +152,94 @@ def test_call_tool_returns_result_with_content(self, cached_genie_call_result):
         assert len(cached_genie_call_result.content) > 0
 
 
+# =============================================================================
+# DBSQL
+# =============================================================================
+
+
+@pytest.mark.integration
+class TestMCPClientDBSQL:
+    """Verify list_tools() and call_tool() against a live DBSQL MCP server."""
+
+    def test_list_tools_returns_expected_tools(self, cached_dbsql_tools_list):
+        tool_names = [t.name for t in cached_dbsql_tools_list]
+        for expected in ["execute_sql", "execute_sql_read_only", "poll_sql_result"]:
+            assert expected in tool_names, f"Expected tool '{expected}' not found in {tool_names}"
+
+    def test_call_tool_execute_sql_read_only(self, dbsql_mcp_client, cached_dbsql_tools_list):
+        """execute_sql_read_only with SHOW CATALOGS should return results."""
+        result = dbsql_mcp_client.call_tool("execute_sql_read_only", {"query": "SHOW CATALOGS"})
+        assert isinstance(result, CallToolResult)
+        assert result.content, "SHOW CATALOGS should return content"
+        assert len(result.content) > 0
+
+
+# =============================================================================
+# Raw streamable_http_client
+# =============================================================================
+
+
+@pytest.mark.integration
+class TestRawStreamableHttpClient:
+    """Verify DatabricksOAuthClientProvider works with the raw MCP SDK streamable_http_client.
+
+    This tests the low-level path: httpx.AsyncClient + DatabricksOAuthClientProvider
+    + streamable_http_client + ClientSession, without going through DatabricksMCPClient.
+    """
+
+    @pytest.mark.asyncio
+    async def test_uc_function_list_and_call(self, uc_function_url, workspace_client):
+        """list_tools + call_tool via raw streamable_http_client for UC functions."""
+        async with raw_mcp_session(uc_function_url, workspace_client) as session:
+            tools_response = await session.list_tools()
+            tools = tools_response.tools
+            assert len(tools) > 0
+            tool_names = [t.name for t in tools]
+            assert any("echo_message" in name for name in tool_names)
+
+            tool_name = next(n for n in tool_names if "echo_message" in n)
+            result = await session.call_tool(tool_name, {"message": "raw_client_test"})
+            assert result.content
+            first = result.content[0]
+            assert hasattr(first, "text")
+            assert "raw_client_test" in str(first.text)
+
+    @pytest.mark.asyncio
+    async def test_vs_list_tools(self, vs_mcp_url, workspace_client):
+        """list_tools via raw streamable_http_client for Vector Search."""
+        async with raw_mcp_session(vs_mcp_url, workspace_client) as session:
+            tools_response = await session.list_tools()
+            assert len(tools_response.tools) > 0
+
+    @pytest.mark.asyncio
+    async def test_dbsql_list_and_call(self, dbsql_mcp_url, workspace_client):
+        """list_tools + call_tool via raw streamable_http_client for DBSQL."""
+        async with raw_mcp_session(dbsql_mcp_url, workspace_client) as session:
+            tools_response = await session.list_tools()
+            tools = tools_response.tools
+            tool_names = [t.name for t in tools]
+            assert "execute_sql_read_only" in tool_names
+
+            result = await session.call_tool("execute_sql_read_only", {"query": "SHOW CATALOGS"})
+            assert result.content
+            assert len(result.content) > 0
+
+    @pytest.mark.asyncio
+    async def test_genie_list_and_call(self, genie_mcp_url, workspace_client):
+        """list_tools + call_tool via raw streamable_http_client for Genie."""
+        async with raw_mcp_session(genie_mcp_url, workspace_client) as session:
+            tools_response = await session.list_tools()
+            tools = tools_response.tools
+            assert len(tools) > 0
+
+            tool = tools[0]
+            properties = tool.inputSchema.get("properties", {})
+            param_name = next(iter(properties), "query")
+            result = await session.call_tool(tool.name, {param_name: "How many rows are there?"})
+            assert result.content
+            assert len(result.content) > 0
+
+
 # =============================================================================
 # Error paths
 # =============================================================================