Simplify migrated_group permission migration fixture (#1136)

## Changes
Applied remaining feedback for #1080

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->
Follow up for #1080

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [x] manually tested
- [ ] verified on staging environment (screenshot attached)

Author: nkvuong

src/databricks/labs/ucx/mixins/fixtures.py

@@ -649,17 +649,12 @@ def make_acc_group(acc, make_random):
 
 
 @pytest.fixture
-def make_migrated_group(acc, ws, make_group, make_acc_group):
+def migrated_group(acc, ws, make_group, make_acc_group):
     """Create a pair of groups in workspace and account. Assign account group to workspace."""
-
-    def inner():
-        ws_group = make_group()
-        acc_group = make_acc_group()
-        acc.workspace_assignment.update(ws.get_workspace_id(), acc_group.id, [iam.WorkspacePermission.USER])
-        # need to return both, as acc_group.id is not in MigratedGroup dataclass
-        return MigratedGroup.partial_info(ws_group, acc_group)
-
-    return inner
+    ws_group = make_group()
+    acc_group = make_acc_group()
+    acc.workspace_assignment.update(ws.get_workspace_id(), acc_group.id, [iam.WorkspacePermission.USER])
+    return MigratedGroup.partial_info(ws_group, acc_group)
 
 
 @pytest.fixture

src/databricks/labs/ucx/workspace_access/groups.py

@@ -99,28 +99,27 @@ def apply_group_permissions_experimental(self, ws: WorkspaceClient) -> bool:
 
     @staticmethod
     def _migrate_group_permissions_paginated(ws: WorkspaceClient, migrated_group: MigratedGroup):
-
         batch_size = 1000
         logger.info(
             f"Migrating permissions from workspace group {migrated_group.name_in_workspace} "
             f"to account group: {migrated_group.name_in_account}."
         )
+        permissions_migrated = 0
         while True:
             response = ws.permission_migration.migrate_permissions(
                 ws.get_workspace_id(),
                 migrated_group.name_in_workspace,
                 migrated_group.name_in_account,
                 size=batch_size,
             )
-            # response shouldn't be empty
-            if response.permissions_migrated is None:
-                break
-            # no more permissions to migrate
-            if response.permissions_migrated == 0:
+            if not response.permissions_migrated:
                 logger.info("No more permission to migrated.")
-                break
-            logger.info(f"Migrated {response.permissions_migrated} permissions.")
-        return 1
+                return permissions_migrated
+            permissions_migrated += response.permissions_migrated
+            logger.info(
+                f"Migrated {response.permissions_migrated} permissions to "
+                f"{migrated_group.name_in_account} account group"
+            )
 
 
 class GroupMigrationStrategy:

tests/integration/workspace_access/test_generic.py

@@ -31,12 +31,11 @@
 def test_instance_pools(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_instance_pool,
     make_instance_pool_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     pool = make_instance_pool()
     make_instance_pool_permissions(
         object_id=pool.instance_pool_id,
@@ -67,13 +66,12 @@ def test_instance_pools(
 def test_clusters(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_cluster,
     make_cluster_permissions,
     env_or_skip,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     cluster = make_cluster(instance_pool_id=env_or_skip("TEST_INSTANCE_POOL_ID"), single_node=True)
     make_cluster_permissions(
         object_id=cluster.cluster_id,
@@ -104,12 +102,11 @@ def test_clusters(
 def test_jobs(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_job,
     make_job_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     job = make_job()
     make_job_permissions(
         object_id=job.job_id,
@@ -140,12 +137,11 @@ def test_jobs(
 def test_pipelines(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_pipeline,
     make_pipeline_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     pipeline = make_pipeline()
     make_pipeline_permissions(
         object_id=pipeline.pipeline_id,
@@ -176,12 +172,11 @@ def test_pipelines(
 def test_cluster_policies(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_cluster_policy,
     make_cluster_policy_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     cluster_policy = make_cluster_policy()
     make_cluster_policy_permissions(
         object_id=cluster_policy.policy_id,
@@ -212,12 +207,11 @@ def test_cluster_policies(
 def test_warehouses(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_warehouse,
     make_warehouse_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     warehouse = make_warehouse()
     make_warehouse_permissions(
         object_id=warehouse.id,
@@ -248,12 +242,11 @@ def test_warehouses(
 def test_models(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_model,
     make_registered_model_permissions,  # pylint: disable=invalid-name
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     model = make_model()
     make_registered_model_permissions(
         object_id=model.id,
@@ -284,12 +277,11 @@ def test_models(
 def test_experiments(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_experiment,
     make_experiment_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     experiment = make_experiment()
     make_experiment_permissions(
         object_id=experiment.experiment_id,
@@ -319,12 +311,11 @@ def test_directories(
     sql_backend,
     inventory_schema,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_directory,
     make_directory_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     directory = make_directory()
     make_directory_permissions(
         object_id=directory,
@@ -364,12 +355,11 @@ def test_notebooks(
     permission_manager: PermissionManager,
     sql_backend,
     inventory_schema,
-    make_migrated_group,
+    migrated_group,
     make_notebook,
     make_notebook_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     notebook = make_notebook()
     make_notebook_permissions(
         object_id=notebook,
@@ -407,11 +397,10 @@ def test_notebooks(
 def test_tokens(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_authorization_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     make_authorization_permissions(
         object_id="tokens",
         permission_level=PermissionLevel.CAN_USE,
@@ -437,7 +426,7 @@ def test_tokens(
 
 
 @retried(on=[BadRequest], timeout=timedelta(minutes=3))
-def test_verify_permissions(ws: WorkspaceClient, make_group, make_migrated_group, make_job, make_job_permissions):
+def test_verify_permissions(ws: WorkspaceClient, make_group, make_job, make_job_permissions):
     group_a = make_group()
     job = make_job()
     make_job_permissions(
@@ -483,12 +472,11 @@ def test_verify_permissions(ws: WorkspaceClient, make_group, make_migrated_group
 def test_endpoints(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_serving_endpoint,
     make_serving_endpoint_permissions,  # pylint: disable=invalid-name
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     endpoint = make_serving_endpoint()
     make_serving_endpoint_permissions(
         object_id=endpoint.response.id,
@@ -513,12 +501,11 @@ def test_endpoints(
 def test_feature_tables(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     make_feature_table,
     make_feature_table_permissions,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     feature_table = make_feature_table()
     make_feature_table_permissions(
         object_id=feature_table["id"],
@@ -545,10 +532,9 @@ def test_feature_tables(
 def test_feature_store_root_page(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
     ws.permissions.update(
         "feature-tables",
         "/root",
@@ -581,10 +567,9 @@ def test_feature_store_root_page(
 def test_models_root_page(
     ws: WorkspaceClient,
     permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
     is_experimental: bool,
 ):
-    migrated_group = make_migrated_group()
 
     ws.permissions.update(
         "registered-models",

tests/integration/workspace_access/test_redash.py

@@ -23,7 +23,7 @@
 def test_permissions_for_redash(
     ws,
     make_group,
-    make_migrated_group,
+    migrated_group,
     make_user,
     make_query,
     make_query_permissions,
@@ -32,7 +32,6 @@ def test_permissions_for_redash(
 ):
     ws_group_temp = make_group()  # simulate temp/backup group
     user = make_user()
-    migrated_group = make_migrated_group()
 
     query = make_query()
     make_query_permissions(

tests/integration/workspace_access/test_scim.py

@@ -27,7 +27,7 @@ def test_some_entitlements(
     ws_group = make_group()
     acc_group = make_acc_group()
     acc.workspace_assignment.update(ws.get_workspace_id(), acc_group.id, [iam.WorkspacePermission.USER])
-    migrated_group = MigratedGroup.partial_info(ws_group, acc_group)
+    migrated_groups = MigratedGroup.partial_info(ws_group, acc_group)
     ws.groups.patch(
         ws_group.id,
         operations=[
@@ -45,9 +45,9 @@ def test_some_entitlements(
     assert "databricks-sql-access" in before
 
     if use_permission_migration_api:
-        MigrationState([migrated_group]).apply_group_permissions_experimental(ws)
+        MigrationState([migrated_groups]).apply_group_permissions_experimental(ws)
     else:
-        apply_tasks(scim_support, [migrated_group])
+        apply_tasks(scim_support, [migrated_groups])
 
     _, after = scim_support.load_for_group(acc_group.id)
     assert "databricks-sql-access" in after

tests/integration/workspace_access/test_secrets.py

@@ -23,13 +23,12 @@
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_permissions_for_secrets(
     ws: WorkspaceClient,
-    make_migrated_group,
+    migrated_group,
     make_secret_scope,
     make_secret_scope_acl,
     permission_manager: PermissionManager,
     use_permission_migration_api: bool,
 ):
-    migrated_group = make_migrated_group()
 
     scope = make_secret_scope()
     make_secret_scope_acl(scope=scope, principal=migrated_group.name_in_workspace, permission=AclPermission.WRITE)

tests/performance/test_performance.py

@@ -5,13 +5,14 @@
 from functools import partial
 from time import process_time
 
+import pytest
 from databricks.labs.blueprint.parallel import Threads
 from databricks.labs.lsql.backends import SqlBackend
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service import iam
 
 from databricks.labs.ucx.workspace_access.base import Permissions
-from databricks.labs.ucx.workspace_access.groups import MigrationState
+from databricks.labs.ucx.workspace_access.groups import MigratedGroup, MigrationState
 from databricks.labs.ucx.workspace_access.manager import PermissionManager
 
 logger = logging.getLogger(__name__)
@@ -27,21 +28,28 @@ class WorkspaceObject:
     type: str
 
 
+@pytest.fixture
+def migrated_group_experimental(acc, ws, make_group, make_acc_group):
+    """Create a pair of groups in workspace and account. Assign account group to workspace."""
+    ws_group = make_group()
+    acc_group = make_acc_group()
+    acc.workspace_assignment.update(ws.get_workspace_id(), acc_group.id, [iam.WorkspacePermission.USER])
+    return MigratedGroup.partial_info(ws_group, acc_group)
+
+
 def test_apply_group_permissions_experimental_performance(
     ws: WorkspaceClient,
     sql_backend: SqlBackend,
     inventory_schema,
-    permission_manager: PermissionManager,
-    make_migrated_group,
+    migrated_group,
+    migrated_group_experimental,
     make_experiment,
     make_model,
     make_cluster_policy,
     env_or_skip,
 ):
     # Making sure this test can only be launched from local
-    env_or_skip("PWD")
-    migrated_group_experimental = make_migrated_group()
-    migrated_group = make_migrated_group()
+    env_or_skip("IDE_PROJECT_ROOTS")
     ws_objects = [
         WorkspaceObject(partial(make_experiment), [iam.PermissionLevel.CAN_MANAGE], "experiment_id", "experiments"),
         WorkspaceObject(
@@ -67,6 +75,7 @@ def test_apply_group_permissions_experimental_performance(
     logger.info(f"Migration using experimental API takes {process_time() - start}s")
 
     start = process_time()
+    permission_manager = PermissionManager.factory(ws, sql_backend, inventory_schema)
     permission_manager.apply_group_permissions(MigrationState([migrated_group]))
     logger.info(f"Migration using normal approach takes {process_time() - start}s")