Tune verify testing (#177)

Author: nfx

src/databricks/labs/ucx/config.py

@@ -117,6 +117,9 @@ class MigrationConfig:
     num_threads: int | None = 4
     log_level: str | None = "INFO"
 
+    # Starting path for notebooks and directories crawler
+    workspace_start_path: str = "/"
+
     def __post_init__(self):
         if self.connect is None:
             self.connect = ConnectConfig()

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

@@ -373,7 +373,7 @@ def inner() -> Iterator[ModelDatabricks]:
 
 class Inventorizers:
     @staticmethod
-    def provide(ws: WorkspaceClient, migration_state: GroupMigrationState, num_threads: int):
+    def provide(ws: WorkspaceClient, migration_state: GroupMigrationState, num_threads: int, start_path: str):
         return [
             RolesAndEntitlementsInventorizer(ws, migration_state),
             TokensAndPasswordsInventorizer(ws),
@@ -434,5 +434,5 @@ def provide(ws: WorkspaceClient, migration_state: GroupMigrationState, num_threa
                 id_attribute="id",
             ),
             SecretScopeInventorizer(ws),
-            WorkspaceInventorizer(ws, num_threads=num_threads),
+            WorkspaceInventorizer(ws, num_threads=num_threads, start_path=start_path),
         ]

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

@@ -301,6 +301,16 @@ 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(
+        self, migration_state: GroupMigrationState, target: Literal["backup", "account"], tuples: list[tuple[str, str]]
+    ):
+        for object_type, object_id in tuples:
+            if object_type == LogicalObjectType.SECRET_SCOPE:
+                self.verify_applied_scope_acls(object_id, migration_state, target)
+            else:
+                self.verify_applied_permissions(object_type, object_id, migration_state, target)
+        self.verify_roles_and_entitlements(migration_state, target)
+
     def verify_applied_permissions(
         self,
         object_type: str,
@@ -336,6 +346,17 @@ def verify_applied_scope_acls(
             dst_permission = self._secret_scope_permission(scope_name, dst_name)
             assert src_permission == dst_permission, "Scope ACLs were not applied correctly"
 
+    def verify_roles_and_entitlements(self, migration_state: GroupMigrationState, target: Literal["backup", "account"]):
+        for el in migration_state.groups:
+            comparison_base = getattr(el, "workspace" if target == "backup" else "backup")
+            comparison_target = getattr(el, target)
+
+            base_group_info = self._ws.groups.get(comparison_base.id)
+            target_group_info = self._ws.groups.get(comparison_target.id)
+
+            assert base_group_info.roles == target_group_info.roles
+            assert base_group_info.entitlements == target_group_info.entitlements
+
     def _secret_scope_permission(self, scope_name: str, group_name: str) -> workspace.AclPermission | None:
         for acl in self._ws.secrets.list_acls(scope=scope_name):
             if acl.principal == group_name:

src/databricks/labs/ucx/toolkits/group_migration.py

@@ -12,6 +12,7 @@
 class GroupMigrationToolkit:
     def __init__(self, config: MigrationConfig):
         self._num_threads = config.num_threads
+        self._workspace_start_path = config.workspace_start_path
 
         databricks_config = config.to_databricks_config()
         self._configure_logger(config.log_level)
@@ -22,9 +23,9 @@ def __init__(self, config: MigrationConfig):
         self._ws.api_client._session.adapters["https://"].max_retries.total = 20
         self._verify_ws_client(self._ws)
 
-        self.group_manager = GroupManager(self._ws, config.groups)
-        self.table_manager = InventoryTableManager(config.inventory, self._ws)
-        self.permissions_manager = PermissionManager(self._ws, self.table_manager)
+        self._group_manager = GroupManager(self._ws, config.groups)
+        self._table_manager = InventoryTableManager(config.inventory, self._ws)
+        self._permissions_manager = PermissionManager(self._ws, self._table_manager)
 
     @staticmethod
     def _verify_ws_client(w: WorkspaceClient):
@@ -40,28 +41,36 @@ def _configure_logger(level: str):
         ucx_logger.setLevel(level)
 
     def prepare_environment(self):
-        self.group_manager.prepare_groups_in_environment()
-        inventorizers = Inventorizers.provide(self._ws, self.group_manager.migration_groups_provider, self._num_threads)
-        self.permissions_manager.set_inventorizers(inventorizers)
+        self._group_manager.prepare_groups_in_environment()
+        inventorizers = Inventorizers.provide(
+            self._ws, self._group_manager.migration_groups_provider, self._num_threads, self._workspace_start_path
+        )
+        self._permissions_manager.set_inventorizers(inventorizers)
 
     def cleanup_inventory_table(self):
-        self.table_manager.cleanup()
+        self._table_manager.cleanup()
 
     def inventorize_permissions(self):
-        self.permissions_manager.inventorize_permissions()
+        self._permissions_manager.inventorize_permissions()
 
     def apply_permissions_to_backup_groups(self):
-        self.permissions_manager.apply_group_permissions(
-            self.group_manager.migration_groups_provider, destination="backup"
+        self._permissions_manager.apply_group_permissions(
+            self._group_manager.migration_groups_provider, destination="backup"
         )
 
+    def verify_permissions_on_backup_groups(self, to_verify):
+        self._permissions_manager.verify(self._group_manager.migration_groups_provider, "backup", to_verify)
+
     def replace_workspace_groups_with_account_groups(self):
-        self.group_manager.replace_workspace_groups_with_account_groups()
+        self._group_manager.replace_workspace_groups_with_account_groups()
 
     def apply_permissions_to_account_groups(self):
-        self.permissions_manager.apply_group_permissions(
-            self.group_manager.migration_groups_provider, destination="account"
+        self._permissions_manager.apply_group_permissions(
+            self._group_manager.migration_groups_provider, destination="account"
         )
 
+    def verify_permissions_on_account_groups(self, to_verify):
+        self._permissions_manager.verify(self._group_manager.migration_groups_provider, "account", to_verify)
+
     def delete_backup_groups(self):
-        self.group_manager.delete_backup_groups()
+        self._group_manager.delete_backup_groups()

tests/integration/conftest.py

@@ -1,4 +1,3 @@
-import json
 import logging
 import os
 import random
@@ -9,24 +8,16 @@
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.core import Config
 
-from databricks.labs.ucx.config import InventoryTable
 from databricks.labs.ucx.providers.mixins.fixtures import *  # noqa: F403
 from databricks.labs.ucx.providers.mixins.sql import StatementExecutionExt
 from databricks.labs.ucx.utils import ThreadedExecution
 
-from .utils import EnvironmentInfo, InstanceProfile
-
 logging.getLogger("tests").setLevel("DEBUG")
 logging.getLogger("databricks.labs.ucx").setLevel("DEBUG")
 
 logger = logging.getLogger(__name__)
 
-NUM_TEST_INSTANCE_PROFILES = int(os.environ.get("NUM_TEST_INSTANCE_PROFILES", 3))
-NUM_TEST_TOKENS = int(os.environ.get("NUM_TEST_TOKENS", 3))
-
-NUM_THREADS = int(os.environ.get("NUM_TEST_THREADS", 20))
-UCX_TESTING_PREFIX = os.environ.get("UCX_TESTING_PREFIX", "ucx")
-Threader = partial(ThreadedExecution, num_threads=NUM_THREADS)
+Threader = partial(ThreadedExecution, num_threads=20)
 load_debug_env_if_runs_from_ide("ucws")  # noqa: F405
 
 
@@ -188,62 +179,3 @@ def inner():
         return ws_group, acc_group
 
     return inner
-
-
-@pytest.fixture
-def env(make_ucx_group, make_random) -> EnvironmentInfo:
-    test_uid = f"ucx_{make_random(4)}"
-    yield EnvironmentInfo(test_uid=test_uid, groups=[make_ucx_group()])
-
-
-@pytest.fixture
-def instance_profiles(env: EnvironmentInfo, ws: WorkspaceClient) -> list[InstanceProfile]:
-    logger.debug("Adding test instance profiles")
-    profiles: list[InstanceProfile] = []
-
-    for i in range(NUM_TEST_INSTANCE_PROFILES):
-        profile_arn = f"arn:aws:iam::123456789:instance-profile/{env.test_uid}-test-{i}"
-        iam_role_arn = f"arn:aws:iam::123456789:role/{env.test_uid}-test-{i}"
-        ws.instance_profiles.add(instance_profile_arn=profile_arn, iam_role_arn=iam_role_arn, skip_validation=True)
-        profiles.append(InstanceProfile(instance_profile_arn=profile_arn, iam_role_arn=iam_role_arn))
-
-    for ws_group, _ in env.groups:
-        if random.choice([True, False]):
-            # randomize to apply roles randomly
-            roles = {
-                "schemas": ["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
-                "Operations": [
-                    {
-                        "op": "add",
-                        "path": "roles",
-                        "value": [{"value": p.instance_profile_arn} for p in random.choices(profiles, k=2)],
-                    }
-                ],
-            }
-            ws.api_client.do("PATCH", f"/api/2.0/preview/scim/v2/Groups/{ws_group.id}", data=json.dumps(roles))
-
-    yield profiles
-
-    logger.debug("Deleting test instance profiles")
-    for profile in profiles:
-        ws.instance_profiles.remove(profile.instance_profile_arn)
-    logger.debug("Test instance profiles deleted")
-
-
-@pytest.fixture()
-def inventory_table(env: EnvironmentInfo, ws: WorkspaceClient, make_catalog, make_schema) -> InventoryTable:
-    catalog, schema = make_schema(make_catalog()).split(".")
-    table = InventoryTable(
-        catalog=catalog,
-        database=schema,
-        name=f"test_inventory_{env.test_uid}",
-    )
-
-    yield table
-
-    logger.debug(f"Cleaning up inventory table {table}")
-    try:
-        ws.tables.delete(table.to_spark())
-        logger.debug(f"Inventory table {table} deleted")
-    except Exception as e:
-        logger.warning(f"Cannot delete inventory table, skipping it. Original exception {e}")