diff --git a/README.md b/README.md index 6252afd..2a99a3d 100644 --- a/README.md +++ b/README.md @@ -542,16 +542,15 @@ See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_ [[back to top](#python-testing-for-databricks)] ### `make_group` fixture -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 [`Group`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/iam.html#databricks.sdk.service.iam.Group). +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 then waiting for up to 3 minutes for the group to be provisioned. +Returns an instance of [`Group`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/iam.html#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: diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index 3eb752c..29faf4a 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -1,14 +1,15 @@ import logging +import warnings from collections.abc import Generator from datetime import timedelta from pytest import fixture +from databricks.sdk import AccountGroupsAPI, GroupsAPI, WorkspaceClient 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.sdk.service.iam import User, Group from databricks.labs.pytester.fixtures.baseline import factory @@ -44,16 +45,15 @@ def create(**kwargs) -> User: @fixture def make_group(ws: WorkspaceClient, make_random, watchdog_purge_suffix): """ - 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`. + 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 then waiting for up to 3 minutes for the 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: @@ -94,7 +94,55 @@ def _scim_values(ids: list[str]) -> list[iam.ComplexValue]: return [iam.ComplexValue(value=x) for x in ids] +def _wait_group_provisioned(interface: AccountGroupsAPI | GroupsAPI, group: Group) -> None: + """Wait for a group to be visible via the supplied group interface. + + Due to consistency issues in the group-management APIs, new groups are not always visible in a consistent manner + after being created or modified. This method can be used to mitigate against this by checking that a group: + + - Is visible via the `.get()` interface; + - Is visible via the `.list()` interface that enumerates groups. + + Visibility is assumed when 2 calls in a row return the expected results. + + Args: + interface: the group-management interface to use for checking whether the groups are visible. + group: the group whose visibility should be verified. + Raises: + NotFound: this is thrown if it takes longer than 90 seconds for the group to become visible via the + management interface. + """ + # Use double-checking to try and compensate for the lack of monotonic consistency with the group-management + # interfaces: two subsequent calls need to succeed for us to proceed. (This is probabilistic, and not a guarantee.) + # The REST API internals cache things for up to 60s, and we see times close to this during tests. The retry timeout + # reflects this: if it's taking much longer then something else is wrong. + group_id = group.id + assert group_id is not None + + @retried(on=[NotFound], timeout=timedelta(seconds=90)) + def _double_get_group() -> None: + interface.get(group_id) + interface.get(group_id) + + def _check_group_in_listing() -> None: + found_groups = interface.list(attributes="id", filter=f'id eq "{group_id}"') + found_ids = {found_group.id for found_group in found_groups} + if group_id not in found_ids: + msg = f"Group id not (yet) found in group listing: {group_id}" + raise NotFound(msg) + + @retried(on=[NotFound], timeout=timedelta(seconds=90)) + def _double_check_group_in_listing() -> None: + _check_group_in_listing() + _check_group_in_listing() + + _double_get_group() + _double_check_group_in_listing() + + def _make_group(name: str, cfg: Config, interface, make_random, watchdog_purge_suffix) -> Generator[Group, None, None]: + _not_specified = object() + @retried(on=[ResourceConflict], timeout=timedelta(seconds=30)) def create( *, @@ -102,7 +150,7 @@ def create( roles: list[str] | None = None, entitlements: list[str] | None = None, display_name: str | None = None, - wait_for_provisioning: bool = False, + wait_for_provisioning: bool | object = _not_specified, **kwargs, ): kwargs["display_name"] = ( @@ -114,6 +162,13 @@ def create( kwargs["roles"] = _scim_values(roles) if entitlements is not None: kwargs["entitlements"] = _scim_values(entitlements) + if wait_for_provisioning is not _not_specified: + warnings.warn( + "Specifying wait_for_provisioning when making a group is deprecated; we always wait.", + DeprecationWarning, + # Call stack is: create()[iam.py], wrapper()[retries.py], inner()[baseline.py], client_code + stacklevel=4, + ) # TODO: REQUEST_LIMIT_EXCEEDED: GetUserPermissionsRequest RPC token bucket limit has been exceeded. group = interface.create(**kwargs) if cfg.is_account_client: @@ -121,12 +176,7 @@ def create( 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() + _wait_group_provisioned(interface, group) return group diff --git a/src/databricks/labs/pytester/fixtures/unwrap.py b/src/databricks/labs/pytester/fixtures/unwrap.py index bc4ad76..c697d60 100644 --- a/src/databricks/labs/pytester/fixtures/unwrap.py +++ b/src/databricks/labs/pytester/fixtures/unwrap.py @@ -65,9 +65,14 @@ def __str__(self): _GENERATORS = set[str]() -def call_stateful(some: Callable[..., Generator[Callable[..., T]]], **kwargs) -> tuple[CallContext, T]: +def call_stateful( + some: Callable[..., Generator[Callable[..., T]]], + call_context_setup: Callable[[CallContext], CallContext] = lambda x: x, + **kwargs, +) -> tuple[CallContext, T]: # pylint: disable=too-complex ctx = CallContext() + ctx = call_context_setup(ctx) drains = [] if len(_GENERATORS) == 0: diff --git a/tests/unit/fixtures/test_iam.py b/tests/unit/fixtures/test_iam.py index f8e46df..f6e1039 100644 --- a/tests/unit/fixtures/test_iam.py +++ b/tests/unit/fixtures/test_iam.py @@ -1,8 +1,15 @@ -from databricks.labs.pytester.fixtures.iam import make_user, make_group, make_acc_group -from databricks.labs.pytester.fixtures.unwrap import call_stateful +import sys +import warnings +from functools import partial +from unittest.mock import call +import pytest -def test_make_user_no_args(): +from databricks.labs.pytester.fixtures.iam import make_acc_group, make_group, make_user, Group +from databricks.labs.pytester.fixtures.unwrap import call_stateful, CallContext + + +def test_make_user_no_args() -> None: ctx, user = call_stateful(make_user) assert ctx is not None assert user is not None @@ -10,17 +17,59 @@ def test_make_user_no_args(): ctx['ws'].users.delete.assert_called_once() -def test_make_group_no_args(): - ctx, group = call_stateful(make_group) - assert ctx is not None +def _setup_groups_api(call_context: CallContext, *, client_fixture_name: str) -> CallContext: + """Minimum mocking of the specific client so that when a group is created it is also visible via the list() method. + This is required because the make_group and make_acc_group fixtures double-check after creating a group to ensure + the group is visible.""" + mock_group = Group(id="an_id") + call_context[client_fixture_name].groups.create.return_value = mock_group + call_context[client_fixture_name].groups.list.return_value = [mock_group] + return call_context + + +def test_make_group_no_args() -> None: + ctx, group = call_stateful(make_group, call_context_setup=partial(_setup_groups_api, client_fixture_name="ws")) + assert group is not None ctx['ws'].groups.create.assert_called_once() + assert ctx['ws'].groups.get.call_args_list == [call("an_id"), call("an_id")] + assert ctx['ws'].groups.list.call_args_list == [ + call(attributes="id", filter='id eq "an_id"'), + call(attributes="id", filter='id eq "an_id"'), + ] ctx['ws'].groups.delete.assert_called_once() -def test_make_acc_group_no_args(): - ctx, group = call_stateful(make_acc_group) - assert ctx is not None +def test_make_acc_group_no_args() -> None: + ctx, group = call_stateful(make_acc_group, call_context_setup=partial(_setup_groups_api, client_fixture_name="acc")) + assert group is not None ctx['acc'].groups.create.assert_called_once() + assert ctx['acc'].groups.get.call_args_list == [call("an_id"), call("an_id")] + assert ctx['acc'].groups.list.call_args_list == [ + call(attributes="id", filter='id eq "an_id"'), + call(attributes="id", filter='id eq "an_id"'), + ] ctx['acc'].groups.delete.assert_called_once() + + +@pytest.mark.parametrize( + "make_group_fixture, client_fixture_name", + [(make_group, "ws"), (make_acc_group, "acc")], +) +def test_make_group_deprecated_arg(make_group_fixture, client_fixture_name) -> None: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + call_stateful( + make_group_fixture, + wait_for_provisioning=True, + call_context_setup=partial(_setup_groups_api, client_fixture_name=client_fixture_name), + ) + + # Check that the expected warning was emitted and attributed to the caller. + (the_warning, *other_warnings) = w + assert not other_warnings + assert issubclass(the_warning.category, DeprecationWarning) + assert "wait_for_provisioning when making a group is deprecated" in str(the_warning.message) + assert the_warning.filename == sys.modules[call_stateful.__module__].__file__