Add retry mechanism to `test_permission_for_files_anonymous_func_migr…

…ation_api` (#3730)

## Changes
Add retry mechanism to
`test_permission_for_files_anonymous_func_migration_api` to account for
eventual consistency

### Linked issues

Resolves #3685

### Tests

- [x] modified integration tests
`test_permission_for_files_anonymous_func_migration_api`

tests/integration/workspace_access/test_tacl.py

@@ -1,7 +1,13 @@
+import datetime as dt
 import json
 import logging
 from collections import defaultdict
+from collections.abc import Callable, Iterable
+from functools import partial
 
+from databricks.sdk.retries import retried
+
+from databricks.labs.ucx.hive_metastore.grants import Grant
 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
@@ -35,29 +41,39 @@ def test_grants_with_permission_migration_api(runtime_ctx, ws, migrated_group, s
     assert {"USAGE", "OWN"} == new_schema_grants["a"][migrated_group.name_in_account]
 
 
-def test_permission_for_files_anonymous_func_migration_api(ws, sql_backend, runtime_ctx, migrated_group):
-    sql_backend.execute(f"GRANT READ_METADATA ON ANY FILE TO `{migrated_group.name_in_workspace}`")
-    sql_backend.execute(f"GRANT SELECT ON ANONYMOUS FUNCTION TO `{migrated_group.name_in_workspace}`")
+def test_permission_for_files_anonymous_func_migration_api(runtime_ctx, migrated_group) -> None:
+    """Test the migration of permissions for any files and anonymous functions."""
+    # TODO: Move migrated group into `runtime_ctx` and follow the `make_` pattern
+    runtime_ctx.sql_backend.execute(f"GRANT READ_METADATA ON ANY FILE TO `{migrated_group.name_in_workspace}`")
+    runtime_ctx.sql_backend.execute(f"GRANT SELECT ON ANONYMOUS FUNCTION TO `{migrated_group.name_in_workspace}`")
 
-    grants = runtime_ctx.grants_crawler
+    MigrationState([migrated_group]).apply_to_groups_with_different_names(runtime_ctx.workspace_client)
 
-    MigrationState([migrated_group]).apply_to_groups_with_different_names(ws)
+    @retried(on=[ValueError], timeout=dt.timedelta(minutes=2))
+    def verify_grants_meet_condition(condition: Callable[[Iterable[Grant]], bool], **kwargs) -> None:
+        """Verify grants meet the condition.
 
-    any_file_actual = {}
-    for any_file_grant in grants.grants(any_file=True):
-        any_file_actual[any_file_grant.principal] = any_file_grant.action_type
+        The method retries the condition check to account for eventual consistency of the permission API.
+        """
+        grants = runtime_ctx.grants_crawler.grants(**kwargs)
+        if not condition(grants):
+            raise ValueError("Grants do not meet condition")
 
-    # both old and new group have permissions
-    assert migrated_group.name_in_workspace not in any_file_actual
-    assert migrated_group.name_in_account in any_file_actual
+    def is_migrated_permission(grants: Iterable[Grant], *, permission: str) -> bool:
+        """Check if the "any file" permissions are migrated.
 
-    anonymous_function_actual = {}
-    for ano_func_grant in grants.grants(anonymous_function=True):
-        anonymous_function_actual[ano_func_grant.principal] = ano_func_grant.action_type
+        The migration permission APIs **move** the permission, thus the group name in workspace should not have the
+        permissions anymore.
+        """
+        principal_actions = {grant.principal: grant.action_type for grant in grants}
+        return (
+            migrated_group.name_in_workspace not in principal_actions
+            and migrated_group.name_in_account in principal_actions
+            and principal_actions[migrated_group.name_in_account] == permission
+        )
 
-    assert migrated_group.name_in_workspace not in anonymous_function_actual
-    assert migrated_group.name_in_account in anonymous_function_actual
-    assert anonymous_function_actual[migrated_group.name_in_account] == "SELECT"
+    verify_grants_meet_condition(partial(is_migrated_permission, permission="READ_METADATA"), any_file=True)
+    verify_grants_meet_condition(partial(is_migrated_permission, permission="SELECT"), anonymous_function=True)
 
 
 def test_permission_for_udfs_migration_api(ws, sql_backend, runtime_ctx, migrated_group):