|
| 1 | +import json |
| 2 | + |
| 3 | +from databricks.sdk import WorkspaceClient |
| 4 | +from databricks.sdk.service import iam |
| 5 | + |
1 | 6 | from databricks.labs.ucx.contexts.workflow_task import RuntimeContext
|
2 | 7 |
|
3 | 8 |
|
4 |
| -def test_fallback_workspace_admin(installation_ctx: RuntimeContext) -> None: |
5 |
| - """Verify that a workspace administrator can be found for our integration environment.""" |
| 9 | +def _find_admins_group_id(ws: WorkspaceClient) -> str: |
| 10 | + for group in ws.groups.list(attributes="id,displayName,meta", filter='displayName eq "admins"'): |
| 11 | + if group.id and group.display_name == "admins" and group.meta and group.meta.resource_type == "WorkspaceGroup": |
| 12 | + return group.id |
| 13 | + msg = f"Could not locate workspace group in {ws.get_workspace_id()}: admins" |
| 14 | + raise RuntimeError(msg) |
| 15 | + |
| 16 | + |
| 17 | +def _find_user_with_name(ws: WorkspaceClient, user_name: str) -> iam.User: |
| 18 | + for user in ws.users.list(attributes="active,groups,roles,userName", filter=f"userName eq {json.dumps(user_name)}"): |
| 19 | + if user.user_name == user_name: |
| 20 | + return user |
| 21 | + # Use debugger if this is not working to avoid internal usernames in public issues or CI logs. |
| 22 | + msg = f"Could not locate user in workspace {ws.get_workspace_id()}: **REDACTED**" |
| 23 | + raise RuntimeError(msg) |
| 24 | + |
| 25 | + |
| 26 | +def _user_is_member_of_group(user: iam.User, group_id: str) -> bool: |
| 27 | + assert user.groups |
| 28 | + return any(g for g in user.groups if g.value == group_id) |
| 29 | + |
| 30 | + |
| 31 | +def _user_has_role(user: iam.User, role_name: str) -> bool: |
| 32 | + assert user.roles |
| 33 | + return any(r for r in user.roles if r.value == role_name) |
| 34 | + |
| 35 | + |
| 36 | +def test_fallback_admin_user(ws, installation_ctx: RuntimeContext) -> None: |
| 37 | + """Verify that an administrator can be found for our integration environment.""" |
6 | 38 | an_admin = installation_ctx.administrator_locator.get_workspace_administrator()
|
7 | 39 |
|
8 |
| - assert "@" in an_admin |
| 40 | + # The specific admin username that we get here depends on the set of current admins in the integration environment, |
| 41 | + # so that can't be checked directly. Instead we check that either: |
| 42 | + # a) they're a member of the 'admins' workspace; or |
| 43 | + # b) are an account admin (with the `account_admin` role assigned). |
| 44 | + # They must also be an active user. |
| 45 | + # |
| 46 | + # References: |
| 47 | + # https://learn.microsoft.com/en-us/azure/databricks/admin/users-groups/groups#account-admin |
| 48 | + # https://learn.microsoft.com/en-us/azure/databricks/admin/users-groups/groups#account-vs-workspace-group |
| 49 | + admins_group_id = _find_admins_group_id(ws) |
| 50 | + the_user = _find_user_with_name(ws, an_admin) |
| 51 | + |
| 52 | + assert an_admin == the_user.user_name and the_user.active |
| 53 | + assert _user_is_member_of_group(the_user, admins_group_id) or _user_has_role(the_user, "account_admin") |
0 commit comments