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 <[email protected]>
databrickslabs · Nov 4, 2024 · 18948fe · 18948fe
1 parent 9d209d2
commit 18948fe
Show file tree

Hide file tree

Showing 5 changed files with 133 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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://[email protected]',
     )
+
+
+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