Improve make_group/make_acc_group fixture consistency (#50)

## Changes

This PR implements some fixture improvements that were originally
implemented downstream:

- When creating an account or workspace group, we now require the group
to be visible via two subsequent `.get()` and `.list()` calls. This
double-check approach mitigates (but doesn't completely eliminate)
problems with consistency of the APIs for working with groups.
- The `wait_for_provisioning` argument to the `make_group` and
`make_acc_group` has now been removed: we always wait. (For
compatibility the argument is still accepted but will trigger a
deprecation warning.)

An incidental change is the internal unit-test plumbing can now be
provided with mock fixtures specific to the test; this is needed to
ensure the double-check implementation can be tested from the unit
tests.

### Linked issues

Supersedes databrickslabs/ucx#2634.

### Tests

- added/updated unit tests
- existing integration tests

Author: asnare

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:

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,15 +94,63 @@ 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(
         *,
         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,
+        wait_for_provisioning: bool | object = _not_specified,
         **kwargs,
     ):
         kwargs["display_name"] = (
@@ -114,19 +162,21 @@ 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:
             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()
+        _wait_group_provisioned(interface, group)
 
         return group
 

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:

tests/unit/fixtures/test_iam.py

@@ -1,26 +1,75 @@
-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
     ctx['ws'].users.create.assert_called_once()
     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__