-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added documentation for
make_group and make_user
- Loading branch information
Showing
6 changed files
with
196 additions
and
120 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,96 +1,111 @@ | ||
| import pytest | ||
| import logging | ||
| from collections.abc import Generator | ||
| from datetime import timedelta | ||
|
|
||
| from pytest import fixture | ||
| from databricks.sdk.config import Config | ||
| from databricks.sdk.errors import ResourceConflict, NotFound | ||
| from databricks.sdk.retries import retried | ||
| from databricks.sdk.service.iam import User, Group | ||
| from databricks.sdk import WorkspaceClient | ||
| from databricks.sdk.service import iam | ||
|
|
||
| from databricks.labs.pytester.fixtures.baseline import factory | ||
| from databricks.labs.pytester.fixtures.baseline import factory, get_purge_suffix | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| @pytest.fixture | ||
| def make_user(ws: WorkspaceClient, make_random): | ||
| """ | ||
| Fixture to manage Databricks workspace users. | ||
| This fixture provides a function to manage Databricks workspace users using the provided workspace (ws). | ||
| Users can be created with a generated user name, and they will be deleted after the test is complete. | ||
| Parameters: | ||
| ----------- | ||
| ws : WorkspaceClient | ||
| A Databricks WorkspaceClient instance. | ||
| make_random : function | ||
| The make_random fixture to generate unique names. | ||
| Returns: | ||
| -------- | ||
| function: | ||
| A function to manage Databricks workspace users. | ||
| Usage Example: | ||
| -------------- | ||
| To manage Databricks workspace users using the make_user fixture: | ||
| .. code-block:: python | ||
|
|
||
| def test_user_management(make_user): | ||
| user_info = make_user() | ||
| assert user_info is not None | ||
| @fixture | ||
| def make_user(ws, make_random, log_workspace_link): | ||
| """ | ||
| This fixture returns a function that creates a Databricks workspace user | ||
| and removes it after the test is complete. In case of random naming conflicts, | ||
| the fixture will retry the creation process for 30 seconds. Returns an instance | ||
| of `databricks.sdk.service.iam.User`. Usage: | ||
| ```python | ||
| def test_new_user(make_user, ws): | ||
| new_user = make_user() | ||
| home_dir = ws.workspace.get_status(f"/Users/{new_user.user_name}") | ||
| assert home_dir.object_type == ObjectType.DIRECTORY | ||
| ``` | ||
| """ | ||
|
|
||
| def create_user(**kwargs): | ||
| return ws.users.create(user_name=f"sdk-{make_random(4)}@example.com".lower(), **kwargs) | ||
|
|
||
| def cleanup_user(user_info): | ||
| ws.users.delete(user_info.id) | ||
|
|
||
| yield from factory("workspace user", create_user, cleanup_user) | ||
| @retried(on=[ResourceConflict], timeout=timedelta(seconds=30)) | ||
| def create(**kwargs) -> User: | ||
| user_name = f"dummy-{make_random(4)}-{get_purge_suffix()}@example.com".lower() | ||
| user = ws.users.create(user_name=user_name, **kwargs) | ||
| log_workspace_link(user.user_name, f'settings/workspace/identity-and-access/users/{user.id}') | ||
| return user | ||
|
|
||
| yield from factory("workspace user", create, lambda item: ws.users.delete(item.id)) | ||
|
|
||
| def _scim_values(ids: list[str]) -> list[iam.ComplexValue]: | ||
| return [iam.ComplexValue(value=x) for x in ids] | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| @fixture | ||
| def make_group(ws: WorkspaceClient, make_random): | ||
| """ | ||
| Fixture to manage Databricks workspace groups. | ||
| This fixture provides a function to manage Databricks workspace groups using the provided workspace (ws). | ||
| Groups can be created with specified members and roles, and they will be deleted after the test is complete. | ||
| Parameters: | ||
| ----------- | ||
| ws : WorkspaceClient | ||
| A Databricks WorkspaceClient instance. | ||
| make_random : function | ||
| The make_random fixture to generate unique names. | ||
| Returns: | ||
| -------- | ||
| function: | ||
| A function to manage Databricks workspace groups. | ||
| This fixture provides a function to manage Databricks workspace groups. Groups can be created with | ||
| specified members and roles, and they will be deleted after the test is complete. Deals with eventual | ||
| consistency issues by retrying the creation process for 30 seconds and allowing up to two minutes | ||
| for group to be provisioned. Returns an instance of `databricks.sdk.service.iam.Group`. | ||
| Keyword arguments: | ||
| * `members` (list of strings): A list of user IDs to add to the group. | ||
| * `roles` (list of strings): A list of roles to assign to the group. | ||
| * `display_name` (str): The display name of the group. | ||
| * `wait_for_provisioning` (bool): If `True`, the function will wait for the group to be provisioned. | ||
| * `entitlements` (list of strings): A list of entitlements to assign to the group. | ||
| The following example creates a group with a single member and independently verifies that the group was created: | ||
| ```python | ||
| def test_new_group(make_group, make_user, ws): | ||
| user = make_user() | ||
| group = make_group(members=[user.id]) | ||
| loaded = ws.groups.get(group.id) | ||
| assert group.display_name == loaded.display_name | ||
| assert group.members == loaded.members | ||
| ``` | ||
| """ | ||
| yield from _make_group("workspace group", ws.config, ws.groups, make_random) | ||
|
|
||
| Usage Example: | ||
| -------------- | ||
| To manage Databricks workspace groups using the make_group fixture: | ||
|
|
||
| .. code-block:: python | ||
| def _scim_values(ids: list[str]) -> list[iam.ComplexValue]: | ||
| return [iam.ComplexValue(value=x) for x in ids] | ||
|
|
||
| def test_group_management(make_group): | ||
| group_info = make_group(members=["user@example.com"], roles=["viewer"]) | ||
| assert group_info is not None | ||
| """ | ||
|
|
||
| def _make_group(name: str, cfg: Config, interface, make_random) -> Generator[Group, None, None]: | ||
| @retried(on=[ResourceConflict], timeout=timedelta(seconds=30)) | ||
| def create( | ||
| *, members: list[str] | None = None, roles: list[str] | None = None, display_name: str | None = None, **kwargs | ||
| *, | ||
| members: list[str] | None = None, | ||
| roles: list[str] | None = None, | ||
| entitlements: list[str] | None = None, | ||
| display_name: str | None = None, | ||
| wait_for_provisioning: bool = False, | ||
| **kwargs, | ||
| ): | ||
| kwargs["display_name"] = f"sdk-{make_random(4)}" if display_name is None else display_name | ||
| kwargs["display_name"] = f"sdk-{make_random(4)}-{get_purge_suffix()}" if display_name is None else display_name | ||
| if members is not None: | ||
| kwargs["members"] = _scim_values(members) | ||
| if roles is not None: | ||
| kwargs["roles"] = _scim_values(roles) | ||
| return ws.groups.create(**kwargs) | ||
| if entitlements is not None: | ||
| kwargs["entitlements"] = _scim_values(entitlements) | ||
| # TODO: REQUEST_LIMIT_EXCEEDED: GetUserPermissionsRequest RPC token bucket limit has been exceeded. | ||
| group = interface.create(**kwargs) | ||
| if cfg.is_account_client: | ||
| logger.info(f"Account group {group.display_name}: {cfg.host}/users/groups/{group.id}/members") | ||
| else: | ||
| logger.info(f"Workspace group {group.display_name}: {cfg.host}#setting/accounts/groups/{group.id}") | ||
|
|
||
| @retried(on=[NotFound], timeout=timedelta(minutes=2)) | ||
| def _wait_for_provisioning() -> None: | ||
| interface.get(group.id) | ||
|
|
||
| if wait_for_provisioning: | ||
| _wait_for_provisioning() | ||
|
|
||
| def cleanup_group(group_info): | ||
| ws.groups.delete(group_info.id) | ||
| return group | ||
|
|
||
| yield from factory("workspace group", create, cleanup_group) | ||
| yield from factory(name, create, lambda item: interface.delete(item.id)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.