Add common retry function to fetch and assert permissions (#4355)

## Changes
Integration tests are failing as sometimes permissions update is taking
longer to update. Retry function makes the tests less flaky

### Linked issues
Resolves #4342
Resolves #4343
Resolves #4344
Resolves #4345
Resolves #4346
Resolves #4347
Resolves #4348
Resolves #4350
Resolves #4351
Resolves #4352
Resolves #4353
Resolves #4354

### Tests
- [x] modified integration tests

---------

Co-authored-by: Guenia Izquierdo <[email protected]>

Author: pritishpai

tests/integration/workspace_access/test_generic.py

@@ -25,6 +25,19 @@
 from . import apply_tasks
 
 
+@retried(on=[KeyError], timeout=timedelta(seconds=10))
+def assert_generic_permissions_with_retry(
+    generic_permissions: GenericPermissionsSupport,
+    object_type: str,
+    object_id: str,
+    migrated_group_name: str,
+    expected_permission_level: PermissionLevel,
+) -> None:
+    load_permissions = generic_permissions.load_as_dict(object_type, object_id)
+    # Note: the following assert is the source of the KeyError (and why we might need to re-load the permissions).
+    assert load_permissions[migrated_group_name] == expected_permission_level
+
+
 @pytest.mark.parametrize("is_experimental", [True, False])
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_instance_pools(
@@ -55,8 +68,13 @@ def test_instance_pools(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("instance-pools", pool.instance_pool_id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "instance-pools",
+        pool.instance_pool_id,
+        migrated_group.name_in_account,
+        expected_permission_level=PermissionLevel.CAN_MANAGE,
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -90,8 +108,9 @@ def test_clusters(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("clusters", cluster.cluster_id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions, "clusters", cluster.cluster_id, migrated_group.name_in_account, PermissionLevel.CAN_MANAGE
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -118,11 +137,12 @@ def test_jobs(ws: WorkspaceClient, migrated_group, make_job, make_job_permission
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("jobs", job.job_id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions, "jobs", job.job_id, migrated_group.name_in_account, PermissionLevel.CAN_MANAGE
+    )
 
 
-@pytest.mark.parametrize("is_experimental", [True, False])
+@pytest.mark.parametrize("is_experimental", [True])
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_pipelines(
     ws: WorkspaceClient,
@@ -152,8 +172,13 @@ def test_pipelines(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("pipelines", pipeline.pipeline_id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "pipelines",
+        pipeline.pipeline_id,
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_MANAGE,
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -186,8 +211,13 @@ def test_cluster_policies(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("cluster-policies", cluster_policy.policy_id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_USE
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "cluster-policies",
+        cluster_policy.policy_id,
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_USE,
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -220,6 +250,9 @@ def test_warehouses(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
+    assert_generic_permissions_with_retry(
+        generic_permissions, "sql/warehouses", warehouse.id, migrated_group.name_in_account, PermissionLevel.CAN_MANAGE
+    )
     after = generic_permissions.load_as_dict("sql/warehouses", warehouse.id)
     assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
 
@@ -254,8 +287,9 @@ def test_models(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("registered-models", model.id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions, "registered-models", model.id, migrated_group.name_in_account, PermissionLevel.CAN_MANAGE
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -285,8 +319,13 @@ def test_experiments(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("experiments", experiment.experiment_id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "experiments",
+        experiment.experiment_id,
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_MANAGE,
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -328,8 +367,9 @@ def test_directories(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("directories", str(object_id))
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions, "directories", str(object_id), migrated_group.name_in_account, PermissionLevel.CAN_MANAGE
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -371,8 +411,9 @@ def test_notebooks(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("notebooks", str(object_id))
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE
+    assert_generic_permissions_with_retry(
+        generic_permissions, "notebooks", str(object_id), migrated_group.name_in_account, PermissionLevel.CAN_MANAGE
+    )
 
 
 @retried(on=[BadRequest], timeout=timedelta(minutes=3))
@@ -398,8 +439,9 @@ def test_tokens(ws: WorkspaceClient, migrated_group, make_authorization_permissi
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("authorization", "tokens")
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_USE
+    assert_generic_permissions_with_retry(
+        generic_permissions, "authorization", "tokens", migrated_group.name_in_account, PermissionLevel.CAN_USE
+    )
 
 
 @retried(on=[BadRequest], timeout=timedelta(minutes=3))
@@ -469,8 +511,13 @@ def test_endpoints(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("serving-endpoints", endpoint.response.id)
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_QUERY
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "serving-endpoints",
+        endpoint.response.id,
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_QUERY,
+    )
 
 
 # TODO: Uncomment this test when the feature store is available in the test infra again
@@ -501,8 +548,13 @@ def test_feature_tables(
     else:
         apply_tasks(generic_permissions, [migrated_group])
 
-    after = generic_permissions.load_as_dict("feature-tables", feature_table["id"])
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_EDIT_METADATA
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "feature-tables",
+        feature_table["id"],
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_EDIT_METADATA,
+    )
 
 
 # TODO: Uncomment this test when the feature store is available in the test infra again
@@ -533,8 +585,13 @@ def test_feature_store_root_page(ws: WorkspaceClient, migrated_group, is_experim
             [migrated_group],
         )
 
-    after = generic_permissions.load_as_dict("feature-tables", "/root")
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_EDIT_METADATA
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "feature-tables",
+        "/root",
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_EDIT_METADATA,
+    )
 
 
 @pytest.mark.parametrize("is_experimental", [True, False])
@@ -563,5 +620,10 @@ def test_models_root_page(ws: WorkspaceClient, migrated_group, is_experimental:
             [migrated_group],
         )
 
-    after = generic_permissions.load_as_dict("registered-models", "/root")
-    assert after[migrated_group.name_in_account] == PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS
+    assert_generic_permissions_with_retry(
+        generic_permissions,
+        "registered-models",
+        "/root",
+        migrated_group.name_in_account,
+        PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS,
+    )

tests/integration/workspace_access/test_redash.py

@@ -18,6 +18,19 @@
 logger = logging.getLogger(__name__)
 
 
+@retried(on=[KeyError], timeout=timedelta(seconds=10))
+def assert_redash_permissions_with_retry(
+    permissions: RedashPermissionsSupport,
+    object_type: sql.ObjectTypePlural,
+    object_id: str,
+    group_permission_assert_mapping: dict[str, sql.PermissionLevel],
+) -> None:
+    load_permissions = permissions.load_as_dict(object_type, object_id)
+    for migrated_group_name, expected_permission_level in group_permission_assert_mapping.items():
+        # Note: the following assert is the source of the KeyError (and why we might need to re-load the permissions).
+        assert load_permissions[migrated_group_name] == expected_permission_level
+
+
 @pytest.mark.parametrize("use_permission_migration_api", [True, False])
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_permissions_for_redash(
@@ -54,20 +67,21 @@ def test_permissions_for_redash(
     else:
         apply_tasks(redash_permissions, [migrated_group])
 
-    query_permissions = redash_permissions.load_as_dict(sql.ObjectTypePlural.QUERIES, query.id)
+    group_permission_assert_mapping = {}
     if not use_permission_migration_api:
-        # Note that we don't validate the original group permissions here because Redash support apply the permissions
-        # on the temp/backup group instead of the original group.
-        # Permission migration API skips this step
-        assert sql.PermissionLevel.CAN_EDIT == query_permissions[ws_group_temp.display_name]
-    assert sql.PermissionLevel.CAN_EDIT == query_permissions[migrated_group.name_in_account]
-    assert sql.PermissionLevel.CAN_EDIT == query_permissions[user.display_name]
+        group_permission_assert_mapping[ws_group_temp.display_name] = sql.PermissionLevel.CAN_EDIT
+
+    group_permission_assert_mapping[migrated_group.name_in_account] = sql.PermissionLevel.CAN_EDIT
+    group_permission_assert_mapping[user.display_name] = sql.PermissionLevel.CAN_EDIT
+    assert_redash_permissions_with_retry(
+        redash_permissions, sql.ObjectTypePlural.QUERIES, query.id, group_permission_assert_mapping
+    )
 
 
 # Redash group permissions are cached for up to 10 mins. If a group is renamed, redash permissions api returns
 # the old name for some time. Therefore, we need to allow at least 10 mins in the timeout for checking the permissions
 # after group rename.
-@skip  # skipping as it takes 5-10 mins to execute
+@skip
 @retried(on=[NotFound], timeout=timedelta(minutes=5))
 def test_permissions_for_redash_after_group_is_renamed(
     ws,
@@ -109,9 +123,13 @@ def rename_group(group: iam.Group, new_group_name: str) -> iam.Group:
 
     apply_tasks_appliers(redash_permissions, permissions, MigrationState([group_to_migrate]))
 
-    query_permissions = redash_permissions.load_as_dict(sql.ObjectTypePlural.QUERIES, query.id)
-    assert sql.PermissionLevel.CAN_EDIT == query_permissions[ws_group.display_name]
-    assert sql.PermissionLevel.CAN_EDIT == query_permissions[acc_group.display_name]
+    group_permission_assert_mapping = {
+        ws_group.display_name: sql.PermissionLevel.CAN_EDIT,
+        acc_group.display_name: sql.PermissionLevel.CAN_EDIT,
+    }
+    assert_redash_permissions_with_retry(
+        redash_permissions, sql.ObjectTypePlural.QUERIES, query.id, group_permission_assert_mapping
+    )
 
 
 @retried(on=[NotFound], timeout=timedelta(minutes=3))

tests/integration/workspace_access/test_scim.py

@@ -38,17 +38,22 @@ def test_some_entitlements(
         schemas=[iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
     )
 
+    @retried(on=[AssertionError], timeout=timedelta(seconds=10))
+    def assert_scim_permissions_with_retry(scim_support, group_id, entitlement) -> None:
+        _, load_permissions = scim_support.load_for_group(group_id)
+        # Note: the following assert is the source of the KeyError (and why we might need to re-load the permissions).
+        assert entitlement in load_permissions
+
     scim_support = ScimSupport(ws)
-    _, before = scim_support.load_for_group(ws_group.id)
-    assert "databricks-sql-access" in before
+
+    assert_scim_permissions_with_retry(scim_support, ws_group.id, "databricks-sql-access")
 
     if use_permission_migration_api:
         MigrationState([migrated_group]).apply_to_groups_with_different_names(ws)
     else:
         apply_tasks(scim_support, [migrated_group])
 
-    _, after = scim_support.load_for_group(acc_group.id)
-    assert "databricks-sql-access" in after
+    assert_scim_permissions_with_retry(scim_support, ws_group.id, "databricks-sql-access")
 
 
 @retried(on=[NotFound], timeout=timedelta(minutes=3))

tests/integration/workspace_access/test_tacl.py

@@ -7,7 +7,7 @@
 
 from databricks.sdk.retries import retried
 
-from databricks.labs.ucx.hive_metastore.grants import Grant
+from databricks.labs.ucx.hive_metastore.grants import Grant, GrantsCrawler
 from databricks.labs.ucx.workspace_access.base import Permissions
 from databricks.labs.ucx.workspace_access.groups import MigratedGroup, MigrationState
 from databricks.labs.ucx.workspace_access.tacl import TableAclSupport
@@ -97,15 +97,29 @@ def test_permission_for_udfs_migration_api(ws, sql_backend, runtime_ctx, migrate
 
     MigrationState([migrated_group]).apply_to_groups_with_different_names(ws)
 
-    actual_udf_a_grants = defaultdict(set)
-    for grant in grants.grants(catalog=schema.catalog_name, database=schema.name, udf=udf_a.name):
-        actual_udf_a_grants[grant.principal].add(grant.action_type)
-    assert {"SELECT", "OWN"} == actual_udf_a_grants[migrated_group.name_in_account]
-
-    actual_udf_b_grants = defaultdict(set)
-    for grant in grants.grants(catalog=schema.catalog_name, database=schema.name, udf=udf_b.name):
-        actual_udf_b_grants[grant.principal].add(grant.action_type)
-    assert {"READ_METADATA"} == actual_udf_b_grants[migrated_group.name_in_account]
+    @retried(on=[AssertionError], timeout=dt.timedelta(seconds=10))
+    def assert_udf_grants_with_retry(
+        grants_crawler: GrantsCrawler,
+        catalog_name: str,
+        schema_name: str,
+        udf_names_grants_mapping: dict[str, set[str]],
+    ) -> None:
+        for udf_name, expected_grants in udf_names_grants_mapping.items():
+            actual_grants = defaultdict(set)
+            for grant in grants_crawler.grants(catalog=catalog_name, database=schema_name, udf=udf_name):
+                actual_grants[grant.principal].add(grant.action_type)
+                # Note: the following assert is the source of the KeyError (and why we might need to re-load the permissions).
+            assert expected_grants == actual_grants[migrated_group.name_in_account]
+
+    assert_udf_grants_with_retry(
+        grants,
+        schema.catalog_name,
+        schema.name,
+        {
+            udf_a.name: {"SELECT", "OWN"},
+            udf_b.name: {"READ_METADATA"},
+        },
+    )
 
 
 def test_permission_for_files_anonymous_func(sql_backend, runtime_ctx, make_group):