Skip to content

Commit efaad3a

Browse files
hectorcast-dbhectorcast-db
authored
Add SPOG integration tests for explicit auth types (#1341)
## Summary Add integration tests that verify workspace and account operations against unified (SPOG) hosts using explicit auth types. These tests ensure that each auth flow works end-to-end when parameters are passed directly to the client constructor, without relying on environment variable auto-detection. ## Motivation Existing integration tests (`test_workspace_operations`, `test_account_operations`) use a pre-built `unified_config` fixture which inherits auth from the environment. This doesn't validate that explicit auth type configuration works correctly against unified hosts. We need per-auth-type coverage to catch regressions in the SPOG auth path for each credential strategy. ## Changes ### New `isolated_env` fixture (`tests/integration/conftest.py`) A factory fixture that takes a debug-env key (for IDE support) and returns a callable for reading env vars. Before returning, it snapshots `os.environ`, then removes every env var that `Config` reads (discovered dynamically from `ConfigAttribute` fields). This ensures `Config` only sees explicitly passed parameters — no env var leakage between auth strategies. Uses `monkeypatch` so the environment is restored after each test. ### `workspace_id` parameter on `WorkspaceClient` (`databricks/sdk/__init__.py`) `WorkspaceClient.__init__` was missing `workspace_id` as an explicit parameter (it was only available via `Config`). Added it to the signature and forwarded it to `Config(...)`. ### SPOG integration tests (`tests/integration/test_unified_profile.py`) | Test | Auth type | Environment | Status | |------|-----------|-------------|--------| | `test_spog_workspace_pat` | PAT | azure-prod-pat | ✅ | | `test_spog_workspace_oauth_m2m` | OAuth M2M | azure-prod-ucws | ✅ | | `test_spog_workspace_azure_client_secret` | Azure Client Secret | azure-prod-ucws | ✅ | | `test_spog_workspace_google_credentials` | Google Credentials | gcp-prod-ucacct | ⏭️ Skipped | | `test_spog_account_oauth_m2m` | OAuth M2M | azure-prod-acct | ✅ | | `test_spog_account_azure_client_secret` | Azure Client Secret | azure-prod-acct | ✅ | | `test_spog_account_google_credentials` | Google Credentials | gcp-prod-ucacct | ✅ | **`test_spog_workspace_google_credentials` is skipped**: The `google-credentials` strategy uses a GCP ID token with `target_audience=cfg.host`. On the unified host, this produces the same token for both account and workspace requests. Account-level APIs accept this token, but workspace-level APIs return 401. Investigation shows the GCP service principal exists in workspace assignments (account-level API) with ADMIN permissions but is not present in `ac.service_principals.list()` — it appears to be a workspace-local SP not federated at the account level. ### Collection hook update (`conftest.py`) Added `unified_config` and `isolated_env` to `pytest_collection_modifyitems` so tests using these fixtures are automatically marked as `integration` and collected by CI (which filters with `-m 'integration'`). ## Test plan - [ ] Verify tests are collected (not silently excluded) in CI runs - [ ] Verify tests pass or skip (if env vars missing) in each environment listed above NO_CHANGELOG=true
1 parent 2fa3228 commit efaad3a

File tree

3 files changed

+197
-1
lines changed

3 files changed

+197
-1
lines changed

‎databricks/sdk/__init__.py‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tests/integration/conftest.py‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import pytest
99

1010
from databricks.sdk import AccountClient, FilesAPI, FilesExt, WorkspaceClient
11+
from databricks.sdk.config import ConfigAttribute
1112
from databricks.sdk.core import Config
1213
from databricks.sdk.environments import Cloud
1314
from databricks.sdk.service.catalog import VolumeType
@@ -46,7 +47,7 @@ def pytest_configure(config):
4647

4748
def pytest_collection_modifyitems(items):
4849
# safer to refer to fixture fns instead of strings
49-
client_fixtures = [x.__name__ for x in [a, w, ucws, ucacct]]
50+
client_fixtures = [x.__name__ for x in [a, w, ucws, ucacct, unified_config, isolated_env]]
5051
for item in items:
5152
current_fixtures = getattr(item, "fixturenames", ())
5253
for requires_client in client_fixtures:
@@ -92,8 +93,10 @@ def unified_config(env_or_skip) -> Config:
9293
_load_debug_env_if_runs_from_ide("account")
9394
env_or_skip("CLOUD_ENV")
9495
config = Config()
96+
config.host = env_or_skip("UNIFIED_HOST")
9597
config.workspace_id = env_or_skip("TEST_WORKSPACE_ID")
9698
config.experimental_is_unified_host = True
99+
config._fix_host_if_needed()
97100
return config
98101

99102

@@ -117,6 +120,39 @@ def ucws(env_or_skip) -> WorkspaceClient:
117120
return WorkspaceClient()
118121

119122

123+
@pytest.fixture
124+
def isolated_env(monkeypatch):
125+
"""Fixture for tests that need to construct a client with explicit parameters
126+
only, without Config picking up env vars automatically.
127+
128+
Usage:
129+
def test_something(isolated_env):
130+
env = isolated_env("workspace") # loads debug-env key when running from IDE
131+
host = env("UNIFIED_HOST")
132+
133+
The first call takes a debug-env key (used only when running from an IDE).
134+
Returns a callable that looks up vars from the original environment snapshot
135+
(skipping the test if missing). All Config-related env vars are removed so
136+
the client only sees explicitly passed parameters."""
137+
138+
def init(debug_env_key: str):
139+
_load_debug_env_if_runs_from_ide(debug_env_key)
140+
original_env = os.environ.copy()
141+
142+
for attr in vars(Config).values():
143+
if isinstance(attr, ConfigAttribute) and attr.env:
144+
monkeypatch.delenv(attr.env, raising=False)
145+
146+
def get(var: str) -> str:
147+
if var not in original_env:
148+
pytest.skip(f"Environment variable {var} is missing")
149+
return original_env[var]
150+
151+
return get
152+
153+
return init
154+
155+
120156
@pytest.fixture(scope="session")
121157
def env_or_skip():
122158

‎tests/integration/test_unified_profile.py‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,17 @@
1+
import pytest
2+
13
from databricks.sdk import AccountClient, WorkspaceClient
4+
from databricks.sdk.environments import Cloud
5+
6+
from .conftest import _is_cloud
27

38

49
def test_workspace_operations(unified_config):
10+
# GCP google-credentials auth produces a workspace-local SP that is not
11+
# federated at the account level. Workspace APIs via the unified host
12+
# return 401 for this SP. See test_spog_workspace_google_credentials.
13+
if _is_cloud(Cloud.GCP):
14+
pytest.skip("google-credentials workspace ops not supported on unified hosts (workspace-local SP)")
515
client = WorkspaceClient(config=unified_config)
616
user = client.current_user.me()
717
assert user is not None
@@ -11,3 +21,151 @@ def test_account_operations(unified_config):
1121
client = AccountClient(config=unified_config)
1222
groups = client.groups.list()
1323
assert groups is not None
24+
25+
26+
# SPOG/W — Workspace operations on unified host with explicit auth
27+
28+
29+
# Environment: azure-prod-pat
30+
def test_spog_workspace_pat(isolated_env):
31+
env = isolated_env("workspace")
32+
host = env("UNIFIED_HOST")
33+
workspace_id = env("TEST_WORKSPACE_ID")
34+
account_id = env("TEST_ACCOUNT_ID")
35+
token = env("DATABRICKS_TOKEN")
36+
ws = WorkspaceClient(
37+
host=host,
38+
workspace_id=workspace_id,
39+
account_id=account_id,
40+
token=token,
41+
)
42+
me = ws.current_user.me()
43+
assert me.user_name
44+
45+
46+
# Environment: azure-prod-ucws
47+
def test_spog_workspace_oauth_m2m(isolated_env):
48+
env = isolated_env("ucws")
49+
host = env("UNIFIED_HOST")
50+
client_id = env("TEST_DATABRICKS_CLIENT_ID")
51+
client_secret = env("TEST_DATABRICKS_CLIENT_SECRET")
52+
workspace_id = env("THIS_WORKSPACE_ID")
53+
account_id = env("TEST_ACCOUNT_ID")
54+
ws = WorkspaceClient(
55+
host=host,
56+
client_id=client_id,
57+
client_secret=client_secret,
58+
workspace_id=workspace_id,
59+
account_id=account_id,
60+
auth_type="oauth-m2m",
61+
)
62+
me = ws.current_user.me()
63+
assert me.user_name
64+
65+
66+
# Environment: azure-prod-ucws
67+
def test_spog_workspace_azure_client_secret(isolated_env):
68+
env = isolated_env("ucws")
69+
host = env("UNIFIED_HOST")
70+
workspace_id = env("THIS_WORKSPACE_ID")
71+
account_id = env("TEST_ACCOUNT_ID")
72+
azure_client_id = env("ARM_CLIENT_ID")
73+
azure_client_secret = env("ARM_CLIENT_SECRET")
74+
azure_tenant_id = env("ARM_TENANT_ID")
75+
ws = WorkspaceClient(
76+
host=host,
77+
workspace_id=workspace_id,
78+
account_id=account_id,
79+
azure_client_id=azure_client_id,
80+
azure_client_secret=azure_client_secret,
81+
azure_tenant_id=azure_tenant_id,
82+
auth_type="azure-client-secret",
83+
)
84+
me = ws.current_user.me()
85+
assert me.user_name
86+
87+
88+
# Environment: gcp-prod-ucacct
89+
def test_spog_workspace_google_credentials(isolated_env):
90+
# google-credentials uses a GCP ID token with target_audience=cfg.host.
91+
# On the unified host this produces the same token for both account and workspace
92+
# requests (identical OIDC exchange, identical audience). Account-level APIs accept
93+
# this token, but workspace-level APIs return 401. The X-Databricks-Org-Id header
94+
# is set correctly. This appears to be a server-side limitation on unified hosts.
95+
pytest.skip("google-credentials ID token is rejected for workspace operations on unified hosts")
96+
env = isolated_env("ucacct")
97+
host = env("UNIFIED_HOST")
98+
account_id = env("DATABRICKS_ACCOUNT_ID")
99+
google_credentials = env("GOOGLE_CREDENTIALS")
100+
google_service_account = env("DATABRICKS_GOOGLE_SERVICE_ACCOUNT")
101+
workspace_id = env("TEST_WORKSPACE_ID")
102+
103+
ws = WorkspaceClient(
104+
host=host,
105+
workspace_id=workspace_id,
106+
account_id=account_id,
107+
google_credentials=google_credentials,
108+
google_service_account=google_service_account,
109+
auth_type="google-credentials",
110+
)
111+
me = ws.current_user.me()
112+
assert me.user_name
113+
114+
115+
# SPOG/A — Account operations on unified host with explicit auth
116+
117+
118+
# Environment: azure-prod-acct
119+
def test_spog_account_oauth_m2m(isolated_env):
120+
env = isolated_env("ucacct")
121+
host = env("UNIFIED_HOST")
122+
account_id = env("DATABRICKS_ACCOUNT_ID")
123+
client_id = env("TEST_DATABRICKS_CLIENT_ID")
124+
client_secret = env("TEST_DATABRICKS_CLIENT_SECRET")
125+
ac = AccountClient(
126+
host=host,
127+
account_id=account_id,
128+
client_id=client_id,
129+
client_secret=client_secret,
130+
auth_type="oauth-m2m",
131+
)
132+
sps = ac.service_principals.list()
133+
next(sps)
134+
135+
136+
# Environment: azure-prod-acct
137+
def test_spog_account_azure_client_secret(isolated_env):
138+
env = isolated_env("ucacct")
139+
host = env("UNIFIED_HOST")
140+
account_id = env("DATABRICKS_ACCOUNT_ID")
141+
azure_client_id = env("ARM_CLIENT_ID")
142+
azure_client_secret = env("ARM_CLIENT_SECRET")
143+
azure_tenant_id = env("ARM_TENANT_ID")
144+
ac = AccountClient(
145+
host=host,
146+
account_id=account_id,
147+
azure_client_id=azure_client_id,
148+
azure_client_secret=azure_client_secret,
149+
azure_tenant_id=azure_tenant_id,
150+
auth_type="azure-client-secret",
151+
)
152+
sps = ac.service_principals.list()
153+
next(sps)
154+
155+
156+
# Environment: gcp-prod-ucacct
157+
def test_spog_account_google_credentials(isolated_env):
158+
env = isolated_env("ucacct")
159+
host = env("UNIFIED_HOST")
160+
account_id = env("DATABRICKS_ACCOUNT_ID")
161+
google_credentials = env("GOOGLE_CREDENTIALS")
162+
google_service_account = env("DATABRICKS_GOOGLE_SERVICE_ACCOUNT")
163+
ac = AccountClient(
164+
host=host,
165+
account_id=account_id,
166+
google_credentials=google_credentials,
167+
google_service_account=google_service_account,
168+
auth_type="google-credentials",
169+
)
170+
sps = ac.service_principals.list()
171+
next(sps)

0 commit comments

Comments
 (0)