From d87403c4900620e8ae5a71c61d033415811a55f9 Mon Sep 17 00:00:00 2001 From: Serge Smertin <259697+nfx@users.noreply.github.com> Date: Fri, 8 Sep 2023 17:04:46 -0700 Subject: [PATCH] Use new fixtures for notebooks and folders (#176) --- .../labs/ucx/providers/mixins/fixtures.py | 6 +- tests/integration/conftest.py | 71 +------------------ tests/integration/test_e2e.py | 62 ++++++++-------- 3 files changed, 38 insertions(+), 101 deletions(-) diff --git a/src/databricks/labs/ucx/providers/mixins/fixtures.py b/src/databricks/labs/ucx/providers/mixins/fixtures.py index 9eab5fa200..dbe3f725b2 100644 --- a/src/databricks/labs/ucx/providers/mixins/fixtures.py +++ b/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)) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index f6bbcf8aae..55159c406a 100644 --- a/tests/integration/conftest.py +++ b/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(".") diff --git a/tests/integration/test_e2e.py b/tests/integration/test_e2e.py index e4e68aca0a..f507dcc1b6 100644 --- a/tests/integration/test_e2e.py +++ b/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")