Remove --experimental-is-unified-host flag and update tests

Remove the two CLI token source tests that asserted the flag was passed,
as we are dropping --experimental-is-unified-host from the CLI args.

Co-authored-by: Isaac

Author: hectorcast-db

databricks/sdk/config.py

@@ -419,6 +419,8 @@ def host_type(self) -> HostType:
         Host type and client type are deprecated. Some hosts can now support both workspace and account APIs.
         This method returns the HostType based on the host pattern, which is not accurate.
         For example, a unified host can support both workspace and account APIs, but WORKSPACE is returned.
+
+        This method still returns the correct value for legacy hosts which only support either workspace or account APIs.
         """
         if not self.host:
             return HostType.WORKSPACE
@@ -436,7 +438,8 @@ def client_type(self) -> ClientType:
         Host type and client type are deprecated. Some hosts can now support both workspace and account APIs.
         This method returns the ClientType based on the host pattern, which is not accurate.
         For example, a unified host can support both workspace and account APIs, but WORKSPACE is returned.
-        Use client_type instead.
+
+        This method still returns the correct value for legacy hosts which only support either workspace or account APIs.
         """
         host_type = self.host_type
 
@@ -452,10 +455,11 @@ def client_type(self) -> ClientType:
     @property
     def is_account_client(self) -> bool:
         """[Deprecated]
-        Host type and client type are deprecated. Clients can now support both workspace and account APIs.
+        Host type and client type are deprecated. Some hosts can now support both workspace and account APIs.
         This method returns True if the host is an accounts host, which is not accurate.
         For example, a unified host can support both workspace and account APIs, but False is returned.
-        Use client_type instead.
+
+        This method still returns the correct value for legacy hosts which only support either workspace or account APIs.
         """
         if not self.host:
             return False

databricks/sdk/credentials_provider.py

@@ -984,14 +984,9 @@ def _validate_token_scopes(self, token: oauth.Token):
     def _build_host_args(cfg: "Config") -> List[str]:
         """Build CLI arguments using --host (legacy path)."""
         args = ["auth", "token", "--host", cfg.host]
-        if cfg.experimental_is_unified_host:
-            # For unified hosts, pass account_id, workspace_id, and experimental flag
-            args += ["--experimental-is-unified-host"]
-            if cfg.account_id:
-                args += ["--account-id", cfg.account_id]
-            if cfg.workspace_id:
-                args += ["--workspace-id", str(cfg.workspace_id)]
-        elif cfg.client_type == ClientType.ACCOUNT:
+        # This is here to support older versions of the Databricks CLI, so we need to keep the client type check.
+        # This won't work for unified hosts, but it is not supposed to.
+        if cfg.client_type == ClientType.ACCOUNT:
             args += ["--account-id", cfg.account_id]
         return args
 

tests/integration/conftest.py

@@ -176,14 +176,10 @@ def _load_debug_env_if_runs_from_ide(key) -> bool:
 
 
 def _is_in_debug() -> bool:
-    return (
-        os.path.basename(sys.argv[0])
-        in [
-            "_jb_pytest_runner.py",
-            "testlauncher.py",
-        ]
-        or "vscode_pytest" in sys.modules
-    )
+    return os.path.basename(sys.argv[0]) in [
+        "_jb_pytest_runner.py",
+        "testlauncher.py",
+    ]
 
 
 def _is_cloud(cloud: Cloud) -> bool:

tests/integration/test_auth.py

@@ -273,6 +273,7 @@ def test_wif_workspace(ucacct, env_or_skip, random):
 
 
 def test_workspace_config_resolves_account_and_workspace_id(w, env_or_skip):
+    """Test that Config resolves account_id and workspace_id from host metadata."""
     env_or_skip("CLOUD_ENV")
 
     config = Config()

tests/test_credentials_provider.py

@@ -310,68 +310,6 @@ def mock_exchange_id_token(id_token: oidc.IdToken):
 class TestDatabricksCliTokenSourceArgs:
     """Tests that DatabricksCliTokenSource constructs correct CLI arguments."""
 
-    def test_unified_host_passes_all_flags(self, mocker):
-        """Unified host should pass --experimental-is-unified-host, --account-id, and --workspace-id."""
-        # Mock the parent class __init__ to capture the command arguments
-        mock_init = mocker.patch.object(
-            credentials_provider.CliTokenSource,
-            "__init__",
-            return_value=None,
-        )
-
-        mock_cfg = Mock()
-        mock_cfg.profile = None
-        mock_cfg.host = "https://example.databricks.com"
-        mock_cfg.experimental_is_unified_host = True
-        mock_cfg.account_id = "test-account-id"
-        mock_cfg.workspace_id = 12345
-        mock_cfg.databricks_cli_path = "/path/to/databricks"
-        mock_cfg.disable_async_token_refresh = False
-
-        credentials_provider.DatabricksCliTokenSource(mock_cfg)
-
-        # Verify the command was constructed correctly
-        call_kwargs = mock_init.call_args
-        cmd = call_kwargs.kwargs["cmd"]
-
-        assert cmd[0] == "/path/to/databricks"
-        assert "auth" in cmd
-        assert "token" in cmd
-        assert "--host" in cmd
-        assert "https://example.databricks.com" in cmd
-        assert "--experimental-is-unified-host" in cmd
-        assert "--account-id" in cmd
-        assert "test-account-id" in cmd
-        assert "--workspace-id" in cmd
-        assert "12345" in cmd
-
-    def test_unified_host_without_workspace_id(self, mocker):
-        """Unified host without workspace_id should only pass --experimental-is-unified-host and --account-id."""
-        mock_init = mocker.patch.object(
-            credentials_provider.CliTokenSource,
-            "__init__",
-            return_value=None,
-        )
-
-        mock_cfg = Mock()
-        mock_cfg.profile = None
-        mock_cfg.host = "https://example.databricks.com"
-        mock_cfg.experimental_is_unified_host = True
-        mock_cfg.account_id = "test-account-id"
-        mock_cfg.workspace_id = None
-        mock_cfg.databricks_cli_path = "/path/to/databricks"
-        mock_cfg.disable_async_token_refresh = False
-
-        credentials_provider.DatabricksCliTokenSource(mock_cfg)
-
-        call_kwargs = mock_init.call_args
-        cmd = call_kwargs.kwargs["cmd"]
-
-        assert "--experimental-is-unified-host" in cmd
-        assert "--account-id" in cmd
-        assert "test-account-id" in cmd
-        assert "--workspace-id" not in cmd
-
     def test_account_client_passes_account_id(self, mocker):
         """Non-unified account client should pass --account-id."""
         mock_init = mocker.patch.object(