Adds retries and uses listing to fetch groups (#4370)

## Changes
The original function being used to fetch groups is taking a long time
to be eventually consistent. Using the same retry logic and fetching the
entire groups list is a much faster approach.

### Linked issues
Resolves #4368
Resolves #4337
Resolved #4334

### Tests
- [x] modified integration tests

Author: pritishpai

tests/integration/conftest.py

@@ -9,7 +9,7 @@
 from dataclasses import replace
 from datetime import timedelta
 from functools import cached_property
-from typing import Literal
+from typing import Literal, NoReturn
 
 import pytest  # pylint: disable=wrong-import-order
 from databricks.labs.blueprint.commands import CommandExecutor
@@ -1465,3 +1465,17 @@ def _run(command: str) -> str:
     except ValueError as err:
         logger.debug(f"pytest_ignore_collect: error: {err}")
         return False
+
+
+@retried(on=[AssertionError], timeout=timedelta(minutes=1))
+def get_group(group_manager: GroupManager, group_name: str) -> NoReturn:
+    """
+    # Group deletion is eventually consistent. Although the group manager tries to wait for convergence, parts of the
+    # API internals have a 60s timeout. As such we should wait at least that long before concluding deletion has not
+    # happened.
+    """
+    deleted_group = group_manager.has_workspace_group(group_name)
+    if deleted_group:
+        logger.info(f"Group {group_name} was not deleted. Retrying...")
+        raise AssertionError(f"Group is not deleted: {group_name}")
+    raise NotFound(f"Group not found: {group_name}")

tests/integration/install/test_installation.py

@@ -4,7 +4,6 @@
 import subprocess
 import sys
 from datetime import timedelta
-from typing import NoReturn
 
 import pytest
 
@@ -28,7 +27,7 @@
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.install import WorkspaceInstaller
 
-from ..conftest import MockInstallationContext
+from ..conftest import MockInstallationContext, get_group
 
 logger = logging.getLogger(__name__)
 
@@ -155,6 +154,8 @@ def test_job_cluster_on_uc_enabled_workpace(ws, installation_ctx) -> None:
 @retried(on=[NotFound, InvalidParameterValue], timeout=timedelta(minutes=5))
 def test_running_real_remove_backup_groups_job(ws: WorkspaceClient, installation_ctx: MockInstallationContext) -> None:
     ws_group_a, _ = installation_ctx.make_ucx_group(wait_for_provisioning=True)
+    assert installation_ctx.group_manager.has_workspace_group(ws_group_a.display_name)
+
     installation_ctx.workspace_installation.run()
     installation_ctx.group_manager.snapshot()
     installation_ctx.group_manager.rename_groups()
@@ -165,17 +166,8 @@ def test_running_real_remove_backup_groups_job(ws: WorkspaceClient, installation
 
     assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
 
-    # Group deletion is eventually consistent. Although the group manager tries to wait for convergence, parts of the
-    # API internals have a 60s timeout. As such we should wait at least that long before concluding deletion has not
-    # happened.
-    # Note: If you are adjusting this, also look at: test_delete_ws_groups_should_delete_renamed_and_reflected_groups_only
-    @retried(on=[KeyError], timeout=timedelta(minutes=5))
-    def get_group(group_id: str) -> NoReturn:
-        ws.groups.get(group_id)
-        raise KeyError(f"Group is not deleted: {group_id}")
-
     with pytest.raises(NotFound):
-        get_group(ws_group_a.id)
+        get_group(installation_ctx.group_manager, ws_group_a.display_name)
 
 
 @retried(on=[NotFound, InvalidParameterValue], timeout=timedelta(minutes=5))

tests/integration/workspace_access/test_groups.py

@@ -2,15 +2,14 @@
 import logging
 import re
 from datetime import timedelta
-from typing import NoReturn
 
 import pytest
 from databricks.sdk.errors import NotFound, ResourceConflict
 from databricks.sdk.retries import retried
 from databricks.sdk.service.iam import Group, ResourceMeta
 
 from databricks.labs.ucx.workspace_access.groups import GroupManager, MigratedGroup
-
+from tests.integration.conftest import get_group
 
 logger = logging.getLogger(__name__)
 
@@ -129,7 +128,7 @@ def test_reflect_account_groups_on_workspace(ws, make_ucx_group, sql_backend, in
     # At this time previous ws level groups aren't deleted
 
 
-@retried(on=[NotFound], timeout=timedelta(minutes=5))
+@retried(on=[NotFound], timeout=timedelta(minutes=2))
 def test_delete_ws_groups_should_delete_renamed_and_reflected_groups_only(
     ws, make_ucx_group, sql_backend, inventory_schema
 ):
@@ -142,24 +141,16 @@ def test_delete_ws_groups_should_delete_renamed_and_reflected_groups_only(
         [ws_group.display_name],
         "ucx-temp-",
     )
+    assert group_manager.has_workspace_group(ws_group.display_name)
     group_manager.rename_groups()
     group_manager.reflect_account_groups_on_workspace()
     group_manager.delete_original_workspace_groups()
 
-    # Group deletion is eventually consistent. Although the group manager tries to wait for convergence, parts of the
-    # API internals have a 60s timeout. As such we should wait at least that long before concluding deletion has not
-    # happened.
-    # Note: If you are adjusting this, also look at: test_running_real_remove_backup_groups_job
-    @retried(on=[KeyError], timeout=timedelta(minutes=5))
-    def get_group(group_id: str) -> NoReturn:
-        ws.groups.get(group_id)
-        raise KeyError(f"Group is not deleted: {group_id}")
-
     with pytest.raises(NotFound):
-        get_group(ws_group.id)
+        get_group(group_manager, ws_group.display_name)
 
 
-@retried(on=[NotFound], timeout=timedelta(minutes=3))
+@retried(on=[NotFound], timeout=timedelta(minutes=2))
 def test_delete_ws_groups_should_not_delete_current_ws_groups(ws, make_ucx_group, sql_backend, inventory_schema):
     ws_group, _ = make_ucx_group()