Use new fixtures for notebooks and folders (#176)

Author: nfx

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

@@ -483,9 +483,9 @@ def create(
         if model_name is None:
             model_name = f"sdk-{make_random(4)}"
 
-        return ws.model_registry.get_model(
-            ws.model_registry.create_model(model_name, **kwargs).registered_model.name
-        ).registered_model_databricks
+        created_model = ws.model_registry.create_model(model_name, **kwargs)
+        model = ws.model_registry.get_model(created_model.registered_model.name)
+        return model.registered_model_databricks
 
     yield from factory("model", create, lambda item: ws.model_registry.delete_model(item.id))
 

tests/integration/conftest.py

@@ -1,4 +1,3 @@
-import io
 import json
 import logging
 import os
@@ -9,16 +8,13 @@
 import pytest
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.core import Config
-from databricks.sdk.service.iam import AccessControlRequest, PermissionLevel
-from databricks.sdk.service.workspace import ObjectInfo, ObjectType
 
 from databricks.labs.ucx.config import InventoryTable
-from databricks.labs.ucx.inventory.types import RequestObjectType
 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, WorkspaceObjects
+from .utils import EnvironmentInfo, InstanceProfile
 
 logging.getLogger("tests").setLevel("DEBUG")
 logging.getLogger("databricks.labs.ucx").setLevel("DEBUG")
@@ -234,71 +230,6 @@ def instance_profiles(env: EnvironmentInfo, ws: WorkspaceClient) -> list[Instanc
     logger.debug("Test instance profiles deleted")
 
 
-@pytest.fixture
-def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObjects:
-    logger.info(f"Creating test workspace objects under /{env.test_uid}")
-    ws.workspace.mkdirs(f"/{env.test_uid}")
-
-    base_dirs = []
-
-    for ws_group, _ in env.groups:
-        _path = f"/{env.test_uid}/{ws_group.display_name}"
-        ws.workspace.mkdirs(_path)
-        object_info = ws.workspace.get_status(_path)
-        base_dirs.append(object_info)
-
-        ws.permissions.set(
-            request_object_type=RequestObjectType.DIRECTORIES,
-            request_object_id=object_info.object_id,
-            access_control_list=[
-                AccessControlRequest(group_name=ws_group.display_name, permission_level=PermissionLevel.CAN_MANAGE)
-            ],
-        )
-
-    notebooks = []
-
-    for nb_idx in range(3):
-        random_group = random.choice([g[0] for g in env.groups])
-        _nb_path = f"/{env.test_uid}/{random_group.display_name}/nb-{nb_idx}.py"
-        ws.workspace.upload(path=_nb_path, content=io.BytesIO(b"print(1)"))
-        # TODO: add a proper test for this
-        # if random.choice([True, False]):
-        #     ws.experiments.create_experiment(name=_nb_path)  # create experiment to test nb-based experiments
-        _nb_obj = ws.workspace.get_status(_nb_path)
-        notebooks.append(_nb_obj)
-        ws.permissions.set(
-            request_object_type=RequestObjectType.NOTEBOOKS,
-            request_object_id=_nb_obj.object_id,
-            access_control_list=[
-                AccessControlRequest(group_name=random_group.display_name, permission_level=PermissionLevel.CAN_EDIT)
-            ],
-        )
-
-    yield WorkspaceObjects(
-        root_dir=ObjectInfo(
-            path=f"/{env.test_uid}",
-            object_type=ObjectType.DIRECTORY,
-            object_id=ws.workspace.get_status(f"/{env.test_uid}").object_id,
-        ),
-        directories=base_dirs,
-        notebooks=notebooks,
-    )
-
-    logger.debug("Deleting test workspace objects")
-    ws.workspace.delete(f"/{env.test_uid}", recursive=True)
-    logger.debug("Test workspace objects deleted")
-
-
-@pytest.fixture
-def verifiable_objects(
-    workspace_objects,
-) -> list[tuple[list, str, RequestObjectType | None]]:
-    _verifiable_objects = [
-        (workspace_objects, "workspace_objects", None),
-    ]
-    yield _verifiable_objects
-
-
 @pytest.fixture()
 def inventory_table(env: EnvironmentInfo, ws: WorkspaceClient, make_catalog, make_schema) -> InventoryTable:
     catalog, schema = make_schema(make_catalog()).split(".")

tests/integration/test_e2e.py

@@ -30,7 +30,6 @@ def _verify_group_permissions(
     objects: list | WorkspaceObjects | None,
     id_attribute: str,
     request_object_type: RequestObjectType | None,
-    ws: WorkspaceClient,
     toolkit: GroupMigrationToolkit,
     target: Literal["backup", "account"],
 ):
@@ -39,30 +38,7 @@ def _verify_group_permissions(
         f"{request_object_type or id_attribute} were applied to {target} groups"
     )
 
-    if id_attribute == "workspace_objects":
-        _workspace_objects: WorkspaceObjects = objects
-
-        # list of groups that source the permissions
-        comparison_base = [
-            getattr(mi, "workspace" if target == "backup" else "backup")
-            for mi in toolkit.group_manager.migration_groups_provider.groups
-        ]
-        # list of groups that are the target of the permissions
-        comparison_target = [getattr(mi, target) for mi in toolkit.group_manager.migration_groups_provider.groups]
-
-        root_permissions = ws.permissions.get(
-            request_object_type=RequestObjectType.DIRECTORIES, request_object_id=_workspace_objects.root_dir.object_id
-        )
-        base_group_names = [g.display_name for g in comparison_base]
-        target_group_names = [g.display_name for g in comparison_target]
-
-        base_acls = [a for a in root_permissions.access_control_list if a.group_name in base_group_names]
-
-        target_acls = [a for a in root_permissions.access_control_list if a.group_name in target_group_names]
-
-        assert len(base_acls) == len(target_acls)
-
-    elif id_attribute == "secret_scopes":
+    if id_attribute == "secret_scopes":
         for scope_name in objects:
             toolkit.permissions_manager.verify_applied_scope_acls(
                 scope_name, toolkit.group_manager.migration_groups_provider, target
@@ -97,7 +73,6 @@ def test_e2e(
     env: EnvironmentInfo,
     inventory_table: InventoryTable,
     ws: WorkspaceClient,
-    verifiable_objects: list[tuple[list, str, RequestObjectType | None]],
     make_instance_pool,
     make_instance_pool_permissions,
     make_cluster,
@@ -110,6 +85,10 @@ def test_e2e(
     make_experiment_permissions,
     make_job,
     make_job_permissions,
+    make_notebook,
+    make_notebook_permissions,
+    make_directory,
+    make_directory_permissions,
     make_pipeline,
     make_pipeline_permissions,
     make_secret_scope,
@@ -121,6 +100,8 @@ def test_e2e(
     logger.debug(f"Test environment: {env.test_uid}")
     ws_group = env.groups[0][0]
 
+    verifiable_objects = []
+
     pool = make_instance_pool()
     make_instance_pool_permissions(
         object_id=pool.instance_pool_id,
@@ -180,6 +161,31 @@ def test_e2e(
         ([experiment], "experiment_id", RequestObjectType.EXPERIMENTS),
     )
 
+    directory = make_directory()
+    make_directory_permissions(
+        object_id=directory,
+        permission_level=random.choice(
+            [PermissionLevel.CAN_READ, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN]
+        ),
+        group_name=ws_group.display_name,
+    )
+
+    verifiable_objects.append(
+        ([ws.workspace.get_status(directory)], "object_id", RequestObjectType.DIRECTORIES),
+    )
+
+    notebook = make_notebook(path=f"{directory}/sample.py")
+    make_notebook_permissions(
+        object_id=notebook,
+        permission_level=random.choice(
+            [PermissionLevel.CAN_READ, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN]
+        ),
+        group_name=ws_group.display_name,
+    )
+    verifiable_objects.append(
+        ([ws.workspace.get_status(notebook)], "object_id", RequestObjectType.NOTEBOOKS),
+    )
+
     job = make_job()
     make_job_permissions(
         object_id=job.job_id,
@@ -255,7 +261,7 @@ def test_e2e(
     toolkit.apply_permissions_to_backup_groups()
 
     for _objects, id_attribute, request_object_type in verifiable_objects:
-        _verify_group_permissions(_objects, id_attribute, request_object_type, ws, toolkit, "backup")
+        _verify_group_permissions(_objects, id_attribute, request_object_type, toolkit, "backup")
 
     _verify_roles_and_entitlements(group_migration_state, ws, "backup")
 
@@ -270,7 +276,7 @@ def test_e2e(
     toolkit.apply_permissions_to_account_groups()
 
     for _objects, id_attribute, request_object_type in verifiable_objects:
-        _verify_group_permissions(_objects, id_attribute, request_object_type, ws, toolkit, "account")
+        _verify_group_permissions(_objects, id_attribute, request_object_type, toolkit, "account")
 
     _verify_roles_and_entitlements(group_migration_state, ws, "account")