Add tests

Author: aravind-segu

databricks/sdk/__init__.py

@@ -1,25 +1,27 @@
 from mcp.client.auth import OAuthClientProvider, TokenStorage
 from mcp.shared.auth import OAuthToken
 
+
 class DatabricksTokenStorage(TokenStorage):
     def __init__(self, config):
         self.config = config
 
-    async def get_tokens(self) -> OAuthToken| None:
+    async def get_tokens(self) -> OAuthToken | None:
         headers = self.config.authenticate()
         token = headers["Authorization"].split("Bearer ")[1]
         return OAuthToken(access_token=token, expires_in=60)
-    
+
+
 class MCP:
     def __init__(self, config):
         self._config = config
         self.databricks_token_storage = DatabricksTokenStorage(config)
 
-    def oauth_provider(self):
+    def get_oauth_provider(self):
         return OAuthClientProvider(
             server_url="",
             client_metadata=None,
             storage=self.databricks_token_storage,
             redirect_handler=None,
             callback_handler=None,
-        )
+        )

databricks/sdk/service/mcp.py

@@ -53,6 +53,8 @@ dev = [
     'langchain-openai; python_version > "3.7"',
     "httpx",
     "build", # some integration tests depend on the databricks-sdk-py wheel
+    "mcp>=1.9.1",
+    "pytest-asyncio"
 ]
 notebook = [
     "ipython>=8,<10",

pyproject.toml

@@ -0,0 +1,20 @@
+import time
+
+import httpx
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_mcp_oauth_provider(monkeypatch):
+    monkeypatch.setattr(time, "time", lambda: 100)
+    from databricks.sdk import WorkspaceClient
+
+    monkeypatch.setenv("DATABRICKS_HOST", "test_host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "test_token")
+
+    w = WorkspaceClient()
+    mcp_oauth_provider = w.mcp.get_oauth_provider()
+
+    request = httpx.Request("GET", "https://example.com")
+    response = await anext(mcp_oauth_provider.async_auth_flow(request))
+    assert response.headers["Authorization"] == "Bearer test_token"