Migrate instance pool to the new fixture framework (#161)

Author: nfx

src/databricks/labs/ucx/inventory/permissions.py

@@ -299,3 +299,27 @@ def apply_group_permissions(self, migration_state: GroupMigrationState, destinat
 
         self._apply_permissions_in_parallel(requests=permission_payloads)
         logger.info(f"All permissions were applied for {destination} groups")
+
+    def verify_applied_permissions(
+        self,
+        object_type: str,
+        object_id: str,
+        migration_state: GroupMigrationState,
+        target: Literal["backup", "account"],
+    ):
+        op = self._ws.permissions.get(object_type, object_id)
+        for info in migration_state.groups:
+            src_permissions = sorted(
+                [_ for _ in op.access_control_list if _.group_name == info.workspace.display_name],
+                key=lambda p: p.group_name,
+            )
+            dst_permissions = sorted(
+                [_ for _ in op.access_control_list if _.group_name == getattr(info, target).display_name],
+                key=lambda p: p.group_name,
+            )
+            assert len(dst_permissions) == len(
+                src_permissions
+            ), f"Target permissions were not applied correctly for {object_type}/{object_id}"
+            assert [t.all_permissions for t in dst_permissions] == [
+                s.all_permissions for s in src_permissions
+            ], f"Target permissions were not applied correctly for {object_type}/{object_id}"

tests/integration/conftest.py

@@ -9,11 +9,7 @@
 import pytest
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.core import Config, DatabricksError
-from databricks.sdk.service.compute import (
-    ClusterDetails,
-    CreateInstancePoolResponse,
-    CreatePolicyResponse,
-)
+from databricks.sdk.service.compute import ClusterDetails, CreatePolicyResponse
 from databricks.sdk.service.iam import AccessControlRequest, PermissionLevel
 from databricks.sdk.service.jobs import CreateResponse
 from databricks.sdk.service.ml import CreateExperimentResponse, ModelDatabricks
@@ -273,31 +269,6 @@ def instance_profiles(env: EnvironmentInfo, ws: WorkspaceClient) -> list[Instanc
     logger.debug("Test instance profiles deleted")
 
 
-@pytest.fixture
-def instance_pools(env: EnvironmentInfo, ws: WorkspaceClient) -> list[CreateInstancePoolResponse]:
-    logger.debug("Creating test instance pools")
-
-    test_instance_pools: list[CreateInstancePoolResponse] = [
-        ws.instance_pools.create(instance_pool_name=f"{env.test_uid}-test-{i}", node_type_id="i3.xlarge")
-        for i in range(NUM_TEST_INSTANCE_POOLS)
-    ]
-
-    _set_random_permissions(
-        test_instance_pools,
-        "instance_pool_id",
-        RequestObjectType.INSTANCE_POOLS,
-        env,
-        ws,
-        permission_levels=[PermissionLevel.CAN_ATTACH_TO, PermissionLevel.CAN_MANAGE],
-    )
-
-    yield test_instance_pools
-
-    logger.debug("Deleting test instance pools")
-    executables = [partial(ws.instance_pools.delete, p.instance_pool_id) for p in test_instance_pools]
-    Threader(executables).run()
-
-
 @pytest.fixture
 def pipelines(env: EnvironmentInfo, ws: WorkspaceClient) -> list[CreatePipelineResponse]:
     logger.debug("Creating test DLT pipelines")
@@ -580,7 +551,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
 
         ws.permissions.set(
             request_object_type=RequestObjectType.DIRECTORIES,
-            request_object_id=object_info._object_id,
+            request_object_id=object_info.object_id,
             access_control_list=[
                 AccessControlRequest(group_name=ws_group.display_name, permission_level=PermissionLevel.CAN_MANAGE)
             ],
@@ -599,7 +570,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
         notebooks.append(_nb_obj)
         ws.permissions.set(
             request_object_type=RequestObjectType.NOTEBOOKS,
-            request_object_id=_nb_obj._object_id,
+            request_object_id=_nb_obj.object_id,
             access_control_list=[
                 AccessControlRequest(group_name=random_group.display_name, permission_level=PermissionLevel.CAN_EDIT)
             ],
@@ -609,7 +580,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
         root_dir=ObjectInfo(
             path=f"/{env.test_uid}",
             object_type=ObjectType.DIRECTORY,
-            object_id=ws.workspace.get_status(f"/{env.test_uid}")._object_id,
+            object_id=ws.workspace.get_status(f"/{env.test_uid}").object_id,
         ),
         directories=base_dirs,
         notebooks=notebooks,
@@ -623,7 +594,6 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
 @pytest.fixture
 def verifiable_objects(
     clusters,
-    instance_pools,
     cluster_policies,
     pipelines,
     jobs,
@@ -639,7 +609,6 @@ def verifiable_objects(
         (secret_scopes, "secret_scopes", None),
         (tokens, "tokens", RequestObjectType.AUTHORIZATION),
         (clusters, "cluster_id", RequestObjectType.CLUSTERS),
-        (instance_pools, "instance_pool_id", RequestObjectType.INSTANCE_POOLS),
         (cluster_policies, "policy_id", RequestObjectType.CLUSTER_POLICIES),
         (pipelines, "pipeline_id", RequestObjectType.PIPELINES),
         (jobs, "job_id", RequestObjectType.JOBS),

tests/integration/test_e2e.py

@@ -1,13 +1,14 @@
 import logging
+import random
 from typing import Literal
 
 import pytest
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.iam import (
     AccessControlRequest,
     AccessControlResponse,
-    ObjectPermissions,
     Permission,
+    PermissionLevel,
 )
 from databricks.sdk.service.workspace import SecretScope
 from pyspark.errors import AnalysisException
@@ -18,6 +19,7 @@
     InventoryConfig,
     InventoryTable,
     MigrationConfig,
+    TaclConfig,
 )
 from databricks.labs.ucx.inventory.types import RequestObjectType
 from databricks.labs.ucx.providers.groups_info import GroupMigrationState
@@ -116,35 +118,12 @@ def _verify_group_permissions(
         assert [p.all_permissions for p in sorted_ws] == [p.all_permissions for p in sorted_target]
     else:
         for _object in objects:
-            _object_permissions: ObjectPermissions = ws.permissions.get(
-                request_object_type, getattr(_object, id_attribute)
+            toolkit.permissions_manager.verify_applied_permissions(
+                request_object_type,
+                getattr(_object, id_attribute),
+                toolkit.group_manager.migration_groups_provider,
+                target,
             )
-            for migration_info in toolkit.group_manager.migration_groups_provider.groups:
-                target_permissions = sorted(
-                    [
-                        p
-                        for p in _object_permissions.access_control_list
-                        if p.group_name == getattr(migration_info, target).display_name
-                    ],
-                    key=lambda p: p.group_name,
-                )
-
-                source_permissions = sorted(
-                    [
-                        p
-                        for p in _object_permissions.access_control_list
-                        if p.group_name == migration_info.workspace.display_name
-                    ],
-                    key=lambda p: p.group_name,
-                )
-
-                assert len(target_permissions) == len(
-                    source_permissions
-                ), f"Target permissions were not applied correctly for object {_object}"
-
-                assert [t.all_permissions for t in target_permissions] == [
-                    s.all_permissions for s in source_permissions
-                ], f"Target permissions were not applied correctly for object {_object}"
 
 
 def _verify_roles_and_entitlements(
@@ -168,14 +147,25 @@ def test_e2e(
     inventory_table: InventoryTable,
     ws: WorkspaceClient,
     verifiable_objects: list[tuple[list, str, RequestObjectType | None]],
+    make_instance_pool,
+    make_instance_pool_permissions,
 ):
     logger.debug(f"Test environment: {env.test_uid}")
+    ws_group = env.groups[0][0]
+
+    pool = make_instance_pool()
+    make_instance_pool_permissions(
+        object_id=pool.instance_pool_id,
+        permission_level=random.choice([PermissionLevel.CAN_ATTACH_TO, PermissionLevel.CAN_MANAGE]),
+        group_name=ws_group.display_name,
+    )
+    verifiable_objects.append(([pool], "instance_pool_id", RequestObjectType.INSTANCE_POOLS))
 
     config = MigrationConfig(
         connect=ConnectConfig.from_databricks_config(ws.config),
-        with_table_acls=False,
         inventory=InventoryConfig(table=inventory_table),
         groups=GroupsConfig(selected=[g[0].display_name for g in env.groups]),
+        tacl=TaclConfig(auto=True),
         log_level="DEBUG",
     )
     toolkit = GroupMigrationToolkit(config)