Skip to content

Commit 18948fe

Browse files
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]>
1 parent 9d209d2 commit 18948fe

File tree

5 files changed

+133
-5
lines changed

5 files changed

+133
-5
lines changed

README.md

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ def test_workspace_operations(ws):
278278
assert len(clusters) >= 0
279279
```
280280

281-
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).
281+
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).
282282

283283

284284
[[back to top](#python-testing-for-databricks)]
@@ -372,7 +372,7 @@ random_string = make_random(k=8)
372372
assert len(random_string) == 8
373373
```
374374

375-
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).
375+
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).
376376

377377

378378
[[back to top](#python-testing-for-databricks)]
@@ -848,7 +848,7 @@ def test_catalog_fixture(make_catalog, make_schema, make_table):
848848
logger.info(f"Created new schema: {from_table_1}")
849849
```
850850

851-
See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
851+
See also [`make_volume`](#make_volume-fixture), [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
852852

853853

854854
[[back to top](#python-testing-for-databricks)]
@@ -870,7 +870,7 @@ def test_catalog_fixture(make_catalog, make_schema, make_table):
870870
logger.info(f"Created new schema: {from_table_1}")
871871
```
872872

873-
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).
873+
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).
874874

875875

876876
[[back to top](#python-testing-for-databricks)]
@@ -932,6 +932,40 @@ def test_storage_credential(env_or_skip, make_storage_credential, make_random):
932932
See also [`ws`](#ws-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
933933

934934

935+
[[back to top](#python-testing-for-databricks)]
936+
937+
### `make_volume` fixture
938+
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).
939+
940+
Keyword Arguments:
941+
* `catalog_name` (str): The name of the catalog where the schema and the volume are.
942+
* `schema_name` (str): The name of the schema where the volume is.
943+
* `name` (str): The name of the volume.
944+
* `comment` (str, optional): The comment attached to the volume.
945+
946+
Usage:
947+
```python
948+
def test_volume_creation(make_catalog, make_schema, make_volume, make_random):
949+
# Create a catalog
950+
catalog = make_catalog()
951+
952+
# Create a schema in the catalog
953+
schema = make_schema(catalog_name=catalog.name)
954+
955+
# Generate a random name for the volume
956+
volume_name = f"dummy_vol_{make_random(6).lower()}"
957+
958+
# Create the volume
959+
volume = make_volume(
960+
catalog_name=catalog.name,
961+
schema_name=schema.name,
962+
name=volume_name
963+
)
964+
```
965+
966+
See also [`ws`](#ws-fixture), [`make_catalog`](#make_catalog-fixture), [`make_schema`](#make_schema-fixture), [`make_random`](#make_random-fixture).
967+
968+
935969
[[back to top](#python-testing-for-databricks)]
936970

937971
### `product_info` fixture

src/databricks/labs/pytester/fixtures/catalog.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
StorageCredentialInfo,
1717
AwsIamRoleRequest,
1818
AzureServicePrincipal,
19+
VolumeInfo,
20+
VolumeType,
1921
)
2022
from databricks.sdk.service.compute import Language
2123
from databricks.labs.pytester.fixtures.baseline import factory
@@ -443,3 +445,71 @@ def remove(storage_credential: StorageCredentialInfo):
443445
ws.storage_credentials.delete(storage_credential.name, force=True)
444446

445447
yield from factory("storage_credential", create, remove)
448+
449+
450+
@fixture
451+
def make_volume(
452+
ws, make_catalog, make_schema, make_random, log_workspace_link
453+
) -> Generator[Callable[..., VolumeInfo], None, None]:
454+
"""
455+
Create a volume and return its info. Remove it after the test. Returns instance of `databricks.sdk.service.catalog.VolumeInfo`.
456+
457+
Keyword Arguments:
458+
* `catalog_name` (str): The name of the catalog where the schema and the volume are.
459+
* `schema_name` (str): The name of the schema where the volume is.
460+
* `name` (str): The name of the volume.
461+
* `comment` (str, optional): The comment attached to the volume.
462+
463+
Usage:
464+
```python
465+
def test_volume_creation(make_catalog, make_schema, make_volume, make_random):
466+
# Create a catalog
467+
catalog = make_catalog()
468+
469+
# Create a schema in the catalog
470+
schema = make_schema(catalog_name=catalog.name)
471+
472+
# Generate a random name for the volume
473+
volume_name = f"dummy_vol_{make_random(6).lower()}"
474+
475+
# Create the volume
476+
volume = make_volume(
477+
catalog_name=catalog.name,
478+
schema_name=schema.name,
479+
name=volume_name
480+
)
481+
```
482+
"""
483+
484+
def create(
485+
*,
486+
catalog_name: str | None = None,
487+
schema_name: str | None = None,
488+
name: str | None = None,
489+
) -> VolumeInfo:
490+
491+
if not catalog_name:
492+
catalog = make_catalog()
493+
catalog_name = catalog.name
494+
495+
if not schema_name:
496+
schema = make_schema(catalog_name=catalog_name)
497+
schema_name = schema.name
498+
499+
if not name:
500+
name = f"dummy_v{make_random(6).lower()}"
501+
502+
volume_info = ws.volumes.create(
503+
catalog_name=catalog_name,
504+
schema_name=schema_name,
505+
name=name,
506+
volume_type=VolumeType.MANAGED,
507+
)
508+
path = f'explore/data/{volume_info.catalog_name}/{volume_info.schema_name}/{volume_info.name}'
509+
log_workspace_link(f'{volume_info.name} volume', path)
510+
return volume_info
511+
512+
def remove(volume_info: VolumeInfo):
513+
ws.volumes.delete(f"{volume_info.catalog_name}.{volume_info.schema_name}.{volume_info.name}")
514+
515+
yield from factory("volume", create, remove)

src/databricks/labs/pytester/fixtures/plugin.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
make_schema,
2424
make_table,
2525
make_storage_credential,
26+
make_volume,
2627
)
2728
from databricks.labs.pytester.fixtures.notebooks import make_directory, make_workspace_file, make_notebook, make_repo
2829
from databricks.labs.pytester.fixtures.permissions import ( # noqa
@@ -96,6 +97,7 @@
9697
'make_schema',
9798
'make_table',
9899
'make_storage_credential',
100+
'make_volume',
99101
'product_info',
100102
'make_model',
101103
'make_experiment',

tests/integration/fixtures/test_catalog.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,10 @@ def test_storage_credential(env_or_skip, make_storage_credential, make_random):
6464
)
6565

6666

67+
def test_make_volume(make_volume):
68+
logger.info(f"Created new volume: {make_volume()}")
69+
70+
6771
def test_remove_after_property_table(ws, make_table, sql_backend):
6872
new_table = make_table()
6973
# TODO: tables.get is currently failing with

tests/unit/fixtures/test_catalog.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from unittest.mock import ANY
22

3-
from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo, SchemaInfo
3+
from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo, SchemaInfo, VolumeType, VolumeInfo
44

55
from databricks.labs.pytester.fixtures.unwrap import call_stateful
66
from databricks.labs.pytester.fixtures.catalog import (
@@ -9,6 +9,7 @@
99
make_catalog,
1010
make_storage_credential,
1111
make_schema,
12+
make_volume,
1213
)
1314

1415

@@ -157,3 +158,20 @@ def test_make_schema() -> None:
157158
full_name='hive_metastore.abc',
158159
storage_location='abfss://[email protected]',
159160
)
161+
162+
163+
def test_make_volume_noargs():
164+
ctx, info = call_stateful(make_volume)
165+
ctx['ws'].volumes.create.assert_called_once()
166+
assert info is not None
167+
168+
169+
def test_make_volume_with_name():
170+
ctx, info = call_stateful(make_volume, name='test_volume')
171+
ctx['ws'].volumes.create.assert_called_once_with(
172+
name='test_volume',
173+
catalog_name="dummy_crandom",
174+
schema_name="dummy_srandom",
175+
volume_type=VolumeType.MANAGED
176+
)
177+
assert info is not None

0 commit comments

Comments
 (0)