Update databricks tests

Author: dbkegley

src/posit/connect/external/databricks.py

@@ -31,7 +31,6 @@
 
 
 POSIT_OAUTH_INTEGRATION_AUTH_TYPE = "posit-oauth-integration"
-POSIT_LOCAL_CLIENT_CREDENTIALS_AUTH_TYPE = "posit-local-client-credentials"
 POSIT_WORKBENCH_AUTH_TYPE = "posit-workbench"
 
 logger = logging.getLogger("posit.sdk")

tests/posit/connect/external/test_databricks.py

@@ -1,4 +1,5 @@
-import base64
+from __future__ import annotations
+
 from unittest.mock import patch
 
 import pytest
@@ -8,27 +9,43 @@
 from posit.connect import Client
 from posit.connect.external.databricks import (
     POSIT_OAUTH_INTEGRATION_AUTH_TYPE,
-    CredentialsProvider,
-    CredentialsStrategy,
-    PositContentCredentialsProvider,
-    PositContentCredentialsStrategy,
-    PositCredentialsProvider,
-    PositCredentialsStrategy,
-    PositLocalContentCredentialsProvider,
-    PositLocalContentCredentialsStrategy,
-    _get_auth_type,
+    POSIT_WORKBENCH_AUTH_TYPE,
+    ConnectStrategy,
+    WorkbenchStrategy,
     _new_bearer_authorization_header,
+    _PositConnectContentCredentialsProvider,
+    _PositConnectViewerCredentialsProvider,
+    databricks_config,
 )
 from posit.connect.oauth import Credentials
 
+try:
+    from databricks.sdk.core import Config, DefaultCredentials
+    from databricks.sdk.credentials_provider import (
+        CredentialsProvider,
+        CredentialsStrategy,
+    )
+
+    # construct a DefaultCredentials CredentialsStrategy
+    # weirdly, you have to call `__call__()` at least once in order to initialize `auth_type()`
+    # This is the expected credentials strategy when none is provided to our databricks_config() helper
+    expected_credentials = DefaultCredentials()  # pyright: ignore[reportPossiblyUnboundVariable]
+    expected_credentials(Config(auth_type="databricks-cli"))  # pyright: ignore[reportPossiblyUnboundVariable]
+
+except ImportError:
+    pytestmark = pytest.mark.skipif(True, reason="requires the Databricks SDK")
+
+
+class mock_strategy(CredentialsStrategy):  # pyright: ignore[reportPossiblyUnboundVariable]
+    def __init__(self, name: str):
+        self.name = name
 
-class mock_strategy(CredentialsStrategy):
     def auth_type(self) -> str:
-        return "local"
+        return self.name
 
-    def __call__(self) -> CredentialsProvider:
+    def __call__(self, *args, **kwargs) -> CredentialsProvider:
         def inner() -> Dict[str, str]:
-            return {"Authorization": "Bearer static-pat-token"}
+            return {"Authorization": f"Bearer {self.name}"}
 
         return inner
 
@@ -84,50 +101,14 @@ def test_new_bearer_authorization_header(self):
         result = _new_bearer_authorization_header(credential)
         assert result == {"Authorization": "Bearer access_token"}
 
-    def test_get_auth_type_local(self):
-        assert _get_auth_type("local-auth") == "local-auth"
-
-    @patch.dict("os.environ", {"RSTUDIO_PRODUCT": "CONNECT"})
-    def test_get_auth_type_connect(self):
-        assert _get_auth_type("local-auth") == POSIT_OAUTH_INTEGRATION_AUTH_TYPE
-
-    @responses.activate
-    def test_local_content_credentials_provider(self):
-        token_url = "https://my-token/url"
-        client_id = "client_id"
-        client_secret = "client_secret_123"
-        basic_auth = f"{client_id}:{client_secret}"
-        b64_basic_auth = base64.b64encode(basic_auth.encode("utf-8")).decode("utf-8")
-
-        responses.post(
-            token_url,
-            match=[
-                responses.matchers.urlencoded_params_matcher(
-                    {
-                        "grant_type": "client_credentials",
-                        "scope": "all-apis",
-                    },
-                ),
-                responses.matchers.header_matcher({"Authorization": f"Basic {b64_basic_auth}"}),
-            ],
-            json={
-                "access_token": "oauth2-m2m-access-token",
-                "token_type": "Bearer",
-                "expires_in": 3600,
-            },
-        )
-
-        cp = PositLocalContentCredentialsProvider(token_url, client_id, client_secret)
-        assert cp() == {"Authorization": "Bearer oauth2-m2m-access-token"}
-
     @patch.dict("os.environ", {"CONNECT_CONTENT_SESSION_TOKEN": "cit"})
     @responses.activate
     def test_posit_content_credentials_provider(self):
         register_mocks()
 
         client = Client(api_key="12345", url="https://connect.example/")
         client._ctx.version = None
-        cp = PositContentCredentialsProvider(client=client)
+        cp = _PositConnectContentCredentialsProvider(client=client)
         assert cp() == {"Authorization": "Bearer content-access-token"}
 
     @responses.activate
@@ -136,95 +117,97 @@ def test_posit_credentials_provider(self):
 
         client = Client(api_key="12345", url="https://connect.example/")
         client._ctx.version = None
-        cp = PositCredentialsProvider(client=client, user_session_token="cit")
+        cp = _PositConnectViewerCredentialsProvider(client=client, user_session_token="cit")
         assert cp() == {"Authorization": "Bearer dynamic-viewer-access-token"}
 
-    @responses.activate
-    def test_local_content_credentials_strategy(self):
-        token_url = "https://my-token/url"
-        client_id = "client_id"
-        client_secret = "client_secret_123"
-        basic_auth = f"{client_id}:{client_secret}"
-        b64_basic_auth = base64.b64encode(basic_auth.encode("utf-8")).decode("utf-8")
-
-        responses.post(
-            token_url,
-            match=[
-                responses.matchers.urlencoded_params_matcher(
-                    {
-                        "grant_type": "client_credentials",
-                        "scope": "all-apis",
-                    },
-                ),
-                responses.matchers.header_matcher({"Authorization": f"Basic {b64_basic_auth}"}),
-            ],
-            json={
-                "access_token": "oauth2-m2m-access-token",
-                "token_type": "Bearer",
-                "expires_in": 3600,
-            },
-        )
+    def test_workbench_strategy(self):
+        # default will attempt to load the workbench profile
+        with pytest.raises(ValueError, match="profile=workbench"):
+            WorkbenchStrategy()
 
-        cs = PositLocalContentCredentialsStrategy(
-            token_url,
-            client_id,
-            client_secret,
+        # providing a Config is allowed
+        cs = WorkbenchStrategy(
+            config=Config(host="https://databricks.com/workspace", token="token")  # pyright: ignore[reportPossiblyUnboundVariable]
         )
+        assert cs.auth_type() == POSIT_WORKBENCH_AUTH_TYPE
         cp = cs()
-        assert cs.auth_type() == "posit-local-client-credentials"
-        assert cp() == {"Authorization": "Bearer oauth2-m2m-access-token"}
 
+        # token from the Config is passed through to the auth header
+        assert cp() == {"Authorization": "Bearer token"}
+
+    @patch.dict("os.environ", {"RSTUDIO_PRODUCT": "CONNECT"})
     @patch.dict("os.environ", {"CONNECT_CONTENT_SESSION_TOKEN": "cit"})
     @responses.activate
-    @patch.dict("os.environ", {"RSTUDIO_PRODUCT": "CONNECT"})
-    def test_posit_content_credentials_strategy(self):
+    def test_connect_strategy(self):
         register_mocks()
-
         client = Client(api_key="12345", url="https://connect.example/")
         client._ctx.version = None
-        cs = PositContentCredentialsStrategy(
-            local_strategy=mock_strategy(),
-            client=client,
-        )
+
+        # the default implementation uses Service Account authentication
+        cs = ConnectStrategy(client=client)
+        assert cs.auth_type() == POSIT_OAUTH_INTEGRATION_AUTH_TYPE
         cp = cs()
-        assert cs.auth_type() == "posit-oauth-integration"
         assert cp() == {"Authorization": "Bearer content-access-token"}
 
-    @responses.activate
-    @patch.dict("os.environ", {"RSTUDIO_PRODUCT": "CONNECT"})
-    def test_posit_credentials_strategy(self):
-        register_mocks()
-
-        client = Client(api_key="12345", url="https://connect.example/")
-        client._ctx.version = None
-        cs = PositCredentialsStrategy(
-            local_strategy=mock_strategy(),
-            user_session_token="cit",
-            client=client,
-        )
+        # if a session token is provided then Viewer auth is used
+        cs = ConnectStrategy(client=client, user_session_token="cit")
         cp = cs()
-        assert cs.auth_type() == "posit-oauth-integration"
         assert cp() == {"Authorization": "Bearer dynamic-viewer-access-token"}
 
-    def test_posit_content_credentials_strategy_fallback(self):
-        # local_strategy is used when the content is running locally
-        client = Client(api_key="12345", url="https://connect.example/")
-        cs = PositContentCredentialsStrategy(
-            local_strategy=mock_strategy(),
-            client=client,
+    def test_databricks_config(self):
+        # credentials_strategy is removed if it is provided
+        cfg = databricks_config(credentials_strategy=mock_strategy("mock"))
+        assert cfg._credentials_strategy is not None
+        assert cfg._credentials_strategy.auth_type() != "mock"
+
+        # kwargs are passed through to the Config() constructor
+        cfg = databricks_config(
+            host="https://databricks.com",
+            cluster_id="cluster_id",
+            warehouse_id="warehouse_id",
+            token="token",
         )
-        cp = cs()
-        assert cs.auth_type() == "local"
-        assert cp() == {"Authorization": "Bearer static-pat-token"}
+        assert cfg.host == "https://databricks.com"
+        assert cfg.cluster_id == "cluster_id"
+        assert cfg.warehouse_id == "warehouse_id"
+        assert cfg.token == "token"
+
+    def test_databricks_config_default(self):
+        cfg = databricks_config(
+            posit_default_strategy=mock_strategy("default"),
+            posit_workbench_strategy=mock_strategy("workbench"),
+            posit_connect_strategy=mock_strategy("connect"),
+        )
+        assert cfg._credentials_strategy.auth_type() == "default"
+
+        # default fallback defaults to DefaultCredentials() when none is provided
+        cfg = databricks_config(auth_type="databricks-cli")
+        assert cfg._credentials_strategy.auth_type() == expected_credentials.auth_type()
+
+    @patch.dict("os.environ", {"RS_SERVER_ADDRESS": "https://workbench.posit.co/"})
+    def test_databricks_config_workbench(self):
+        cfg = databricks_config(
+            posit_default_strategy=mock_strategy("default"),
+            posit_workbench_strategy=mock_strategy("workbench"),
+            posit_connect_strategy=mock_strategy("connect"),
+        )
+        assert cfg._credentials_strategy.auth_type() == "workbench"
 
-    def test_posit_credentials_strategy_fallback(self):
-        # local_strategy is used when the content is running locally
-        client = Client(api_key="12345", url="https://connect.example/")
-        cs = PositCredentialsStrategy(
-            local_strategy=mock_strategy(),
-            user_session_token="cit",
-            client=client,
+        # workbench defaults to DefaultCredentials() when none is provided
+        cfg = databricks_config(auth_type="databricks-cli")
+        assert cfg._credentials_strategy.auth_type() == expected_credentials.auth_type()
+
+    @patch.dict("os.environ", {"CONNECT_API_KEY": "API_KEY"})
+    @patch.dict("os.environ", {"CONNECT_SERVER": "https://connect.posit.co/"})
+    @patch.dict("os.environ", {"RSTUDIO_PRODUCT": "CONNECT"})
+    def test_databricks_config_connect(self):
+        cfg = databricks_config(
+            posit_default_strategy=mock_strategy("default"),
+            posit_workbench_strategy=mock_strategy("workbench"),
+            posit_connect_strategy=mock_strategy("connect"),
         )
-        cp = cs()
-        assert cs.auth_type() == "local"
-        assert cp() == {"Authorization": "Bearer static-pat-token"}
+        assert cfg._credentials_strategy.auth_type() == "connect"
+
+        # connect defaults to ConnectStrategy() when none is provided
+        cfg = databricks_config()
+        assert cfg._credentials_strategy.auth_type() == ConnectStrategy().auth_type()