From 18948fed95b453757af6f244dcf90a9786e83775 Mon Sep 17 00:00:00 2001 From: Christopher Grant Date: Mon, 4 Nov 2024 09:14:10 -0800 Subject: [PATCH] Add Volume Fixture (#72) ## Changes This PR adds a Managed Volume fixture from Unity Catalog, allowing testers to create and use a random volume in the catalog. ### Linked issues Resolves #70 . ### Tests - [x] manually tested - [x] added unit tests - [x] added integration tests - [ ] verified on staging environment (screenshot attached) --------- Co-authored-by: chris.grant --- README.md | 42 +++++++++-- .../labs/pytester/fixtures/catalog.py | 70 +++++++++++++++++++ .../labs/pytester/fixtures/plugin.py | 2 + tests/integration/fixtures/test_catalog.py | 4 ++ tests/unit/fixtures/test_catalog.py | 20 +++++- 5 files changed, 133 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index f904470..d209141 100644 --- a/README.md +++ b/README.md @@ -278,7 +278,7 @@ def test_workspace_operations(ws): assert len(clusters) >= 0 ``` -See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture). +See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file`](#make_workspace_file-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture). [[back to top](#python-testing-for-databricks)] @@ -372,7 +372,7 @@ random_string = make_random(k=8) assert len(random_string) == 8 ``` -See also [`make_acc_group`](#make_acc_group-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_workspace_file`](#make_workspace_file-fixture). +See also [`make_acc_group`](#make_acc_group-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_feature_table`](#make_feature_table-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_volume`](#make_volume-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_workspace_file`](#make_workspace_file-fixture). [[back to top](#python-testing-for-databricks)] @@ -848,7 +848,7 @@ def test_catalog_fixture(make_catalog, make_schema, make_table): logger.info(f"Created new schema: {from_table_1}") ``` -See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture). +See also [`make_volume`](#make_volume-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture). [[back to top](#python-testing-for-databricks)] @@ -870,7 +870,7 @@ def test_catalog_fixture(make_catalog, make_schema, make_table): logger.info(f"Created new schema: {from_table_1}") ``` -See also [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`sql_backend`](#sql_backend-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture). +See also [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_volume`](#make_volume-fixture), [`sql_backend`](#sql_backend-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture). [[back to top](#python-testing-for-databricks)] @@ -932,6 +932,40 @@ def test_storage_credential(env_or_skip, make_storage_credential, make_random): See also [`ws`](#ws-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture). +[[back to top](#python-testing-for-databricks)] + +### `make_volume` fixture +Create a volume and return its info. Remove it after the test. Returns instance of [`VolumeInfo`](https://databricks-sdk-py.readthedocs.io/en/latest/dbdataclasses/catalog.html#databricks.sdk.service.catalog.VolumeInfo). + +Keyword Arguments: +* `catalog_name` (str): The name of the catalog where the schema and the volume are. +* `schema_name` (str): The name of the schema where the volume is. +* `name` (str): The name of the volume. +* `comment` (str, optional): The comment attached to the volume. + +Usage: +```python +def test_volume_creation(make_catalog, make_schema, make_volume, make_random): + # Create a catalog + catalog = make_catalog() + + # Create a schema in the catalog + schema = make_schema(catalog_name=catalog.name) + + # Generate a random name for the volume + volume_name = f"dummy_vol_{make_random(6).lower()}" + + # Create the volume + volume = make_volume( + catalog_name=catalog.name, + schema_name=schema.name, + name=volume_name + ) +``` + +See also [`ws`](#ws-fixture), [`make_catalog`](#make_catalog-fixture), [`make_schema`](#make_schema-fixture), [`make_random`](#make_random-fixture). + + [[back to top](#python-testing-for-databricks)] ### `product_info` fixture diff --git a/src/databricks/labs/pytester/fixtures/catalog.py b/src/databricks/labs/pytester/fixtures/catalog.py index f1d2f0a..e76bc73 100644 --- a/src/databricks/labs/pytester/fixtures/catalog.py +++ b/src/databricks/labs/pytester/fixtures/catalog.py @@ -16,6 +16,8 @@ StorageCredentialInfo, AwsIamRoleRequest, AzureServicePrincipal, + VolumeInfo, + VolumeType, ) from databricks.sdk.service.compute import Language from databricks.labs.pytester.fixtures.baseline import factory @@ -443,3 +445,71 @@ def remove(storage_credential: StorageCredentialInfo): ws.storage_credentials.delete(storage_credential.name, force=True) yield from factory("storage_credential", create, remove) + + +@fixture +def make_volume( + ws, make_catalog, make_schema, make_random, log_workspace_link +) -> Generator[Callable[..., VolumeInfo], None, None]: + """ + Create a volume and return its info. Remove it after the test. Returns instance of `databricks.sdk.service.catalog.VolumeInfo`. + + Keyword Arguments: + * `catalog_name` (str): The name of the catalog where the schema and the volume are. + * `schema_name` (str): The name of the schema where the volume is. + * `name` (str): The name of the volume. + * `comment` (str, optional): The comment attached to the volume. + + Usage: + ```python + def test_volume_creation(make_catalog, make_schema, make_volume, make_random): + # Create a catalog + catalog = make_catalog() + + # Create a schema in the catalog + schema = make_schema(catalog_name=catalog.name) + + # Generate a random name for the volume + volume_name = f"dummy_vol_{make_random(6).lower()}" + + # Create the volume + volume = make_volume( + catalog_name=catalog.name, + schema_name=schema.name, + name=volume_name + ) + ``` + """ + + def create( + *, + catalog_name: str | None = None, + schema_name: str | None = None, + name: str | None = None, + ) -> VolumeInfo: + + if not catalog_name: + catalog = make_catalog() + catalog_name = catalog.name + + if not schema_name: + schema = make_schema(catalog_name=catalog_name) + schema_name = schema.name + + if not name: + name = f"dummy_v{make_random(6).lower()}" + + volume_info = ws.volumes.create( + catalog_name=catalog_name, + schema_name=schema_name, + name=name, + volume_type=VolumeType.MANAGED, + ) + path = f'explore/data/{volume_info.catalog_name}/{volume_info.schema_name}/{volume_info.name}' + log_workspace_link(f'{volume_info.name} volume', path) + return volume_info + + def remove(volume_info: VolumeInfo): + ws.volumes.delete(f"{volume_info.catalog_name}.{volume_info.schema_name}.{volume_info.name}") + + yield from factory("volume", create, remove) diff --git a/src/databricks/labs/pytester/fixtures/plugin.py b/src/databricks/labs/pytester/fixtures/plugin.py index ee5abfe..20097ce 100644 --- a/src/databricks/labs/pytester/fixtures/plugin.py +++ b/src/databricks/labs/pytester/fixtures/plugin.py @@ -23,6 +23,7 @@ make_schema, make_table, make_storage_credential, + make_volume, ) from databricks.labs.pytester.fixtures.notebooks import make_directory, make_workspace_file, make_notebook, make_repo from databricks.labs.pytester.fixtures.permissions import ( # noqa @@ -96,6 +97,7 @@ 'make_schema', 'make_table', 'make_storage_credential', + 'make_volume', 'product_info', 'make_model', 'make_experiment', diff --git a/tests/integration/fixtures/test_catalog.py b/tests/integration/fixtures/test_catalog.py index 5d5f6cc..957660f 100644 --- a/tests/integration/fixtures/test_catalog.py +++ b/tests/integration/fixtures/test_catalog.py @@ -64,6 +64,10 @@ def test_storage_credential(env_or_skip, make_storage_credential, make_random): ) +def test_make_volume(make_volume): + logger.info(f"Created new volume: {make_volume()}") + + def test_remove_after_property_table(ws, make_table, sql_backend): new_table = make_table() # TODO: tables.get is currently failing with diff --git a/tests/unit/fixtures/test_catalog.py b/tests/unit/fixtures/test_catalog.py index 82e8023..64049d7 100644 --- a/tests/unit/fixtures/test_catalog.py +++ b/tests/unit/fixtures/test_catalog.py @@ -1,6 +1,6 @@ from unittest.mock import ANY -from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo, SchemaInfo +from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo, SchemaInfo, VolumeType, VolumeInfo from databricks.labs.pytester.fixtures.unwrap import call_stateful from databricks.labs.pytester.fixtures.catalog import ( @@ -9,6 +9,7 @@ make_catalog, make_storage_credential, make_schema, + make_volume, ) @@ -157,3 +158,20 @@ def test_make_schema() -> None: full_name='hive_metastore.abc', storage_location='abfss://container1@storage.com', ) + + +def test_make_volume_noargs(): + ctx, info = call_stateful(make_volume) + ctx['ws'].volumes.create.assert_called_once() + assert info is not None + + +def test_make_volume_with_name(): + ctx, info = call_stateful(make_volume, name='test_volume') + ctx['ws'].volumes.create.assert_called_once_with( + name='test_volume', + catalog_name="dummy_crandom", + schema_name="dummy_srandom", + volume_type=VolumeType.MANAGED + ) + assert info is not None