Skip to content

Commit 864fdff

Browse files
authored
Add retry mechanism to test_permission_for_files_anonymous_func_migration_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`
1 parent 05593cd commit 864fdff

File tree

1 file changed

+33
-17
lines changed

1 file changed

+33
-17
lines changed

tests/integration/workspace_access/test_tacl.py

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,13 @@
1+
import datetime as dt
12
import json
23
import logging
34
from collections import defaultdict
5+
from collections.abc import Callable, Iterable
6+
from functools import partial
47

8+
from databricks.sdk.retries import retried
9+
10+
from databricks.labs.ucx.hive_metastore.grants import Grant
511
from databricks.labs.ucx.workspace_access.base import Permissions
612
from databricks.labs.ucx.workspace_access.groups import MigratedGroup, MigrationState
713
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
3541
assert {"USAGE", "OWN"} == new_schema_grants["a"][migrated_group.name_in_account]
3642

3743

38-
def test_permission_for_files_anonymous_func_migration_api(ws, sql_backend, runtime_ctx, migrated_group):
39-
sql_backend.execute(f"GRANT READ_METADATA ON ANY FILE TO `{migrated_group.name_in_workspace}`")
40-
sql_backend.execute(f"GRANT SELECT ON ANONYMOUS FUNCTION TO `{migrated_group.name_in_workspace}`")
44+
def test_permission_for_files_anonymous_func_migration_api(runtime_ctx, migrated_group) -> None:
45+
"""Test the migration of permissions for any files and anonymous functions."""
46+
# TODO: Move migrated group into `runtime_ctx` and follow the `make_` pattern
47+
runtime_ctx.sql_backend.execute(f"GRANT READ_METADATA ON ANY FILE TO `{migrated_group.name_in_workspace}`")
48+
runtime_ctx.sql_backend.execute(f"GRANT SELECT ON ANONYMOUS FUNCTION TO `{migrated_group.name_in_workspace}`")
4149

42-
grants = runtime_ctx.grants_crawler
50+
MigrationState([migrated_group]).apply_to_groups_with_different_names(runtime_ctx.workspace_client)
4351

44-
MigrationState([migrated_group]).apply_to_groups_with_different_names(ws)
52+
@retried(on=[ValueError], timeout=dt.timedelta(minutes=2))
53+
def verify_grants_meet_condition(condition: Callable[[Iterable[Grant]], bool], **kwargs) -> None:
54+
"""Verify grants meet the condition.
4555
46-
any_file_actual = {}
47-
for any_file_grant in grants.grants(any_file=True):
48-
any_file_actual[any_file_grant.principal] = any_file_grant.action_type
56+
The method retries the condition check to account for eventual consistency of the permission API.
57+
"""
58+
grants = runtime_ctx.grants_crawler.grants(**kwargs)
59+
if not condition(grants):
60+
raise ValueError("Grants do not meet condition")
4961

50-
# both old and new group have permissions
51-
assert migrated_group.name_in_workspace not in any_file_actual
52-
assert migrated_group.name_in_account in any_file_actual
62+
def is_migrated_permission(grants: Iterable[Grant], *, permission: str) -> bool:
63+
"""Check if the "any file" permissions are migrated.
5364
54-
anonymous_function_actual = {}
55-
for ano_func_grant in grants.grants(anonymous_function=True):
56-
anonymous_function_actual[ano_func_grant.principal] = ano_func_grant.action_type
65+
The migration permission APIs **move** the permission, thus the group name in workspace should not have the
66+
permissions anymore.
67+
"""
68+
principal_actions = {grant.principal: grant.action_type for grant in grants}
69+
return (
70+
migrated_group.name_in_workspace not in principal_actions
71+
and migrated_group.name_in_account in principal_actions
72+
and principal_actions[migrated_group.name_in_account] == permission
73+
)
5774

58-
assert migrated_group.name_in_workspace not in anonymous_function_actual
59-
assert migrated_group.name_in_account in anonymous_function_actual
60-
assert anonymous_function_actual[migrated_group.name_in_account] == "SELECT"
75+
verify_grants_meet_condition(partial(is_migrated_permission, permission="READ_METADATA"), any_file=True)
76+
verify_grants_meet_condition(partial(is_migrated_permission, permission="SELECT"), anonymous_function=True)
6177

6278

6379
def test_permission_for_udfs_migration_api(ws, sql_backend, runtime_ctx, migrated_group):

0 commit comments

Comments
 (0)