Updating make_schema fixture to include location to create managed schema (#66)

<!-- REMOVE IRRELEVANT COMMENTS BEFORE CREATING A PULL REQUEST -->
## Changes
Added option to pass location field in make_schema fixture to create a
managed schema


Resolves #65 

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [ ] added unit tests
- [ ] added integration tests

Author: HariGS-DB

README.md

@@ -859,6 +859,7 @@ Create a schema and return its info. Remove it after the test. Returns instance
 Keyword Arguments:
 * `catalog_name` (str): The name of the catalog where the schema will be created. Default is `hive_metastore`.
 * `name` (str): The name of the schema. Default is a random string.
+* `location` (str): The path to the location if it should be a managed schema.
 
 Usage:
 ```python

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

@@ -241,17 +241,15 @@ def remove(table_info: TableInfo):
 
 @fixture
 def make_schema(
-    sql_backend,
-    make_random,
-    log_workspace_link,
-    watchdog_remove_after,
+    sql_backend, make_random, log_workspace_link, watchdog_remove_after
 ) -> Generator[Callable[..., SchemaInfo], None, None]:
     """
     Create a schema and return its info. Remove it after the test. Returns instance of `databricks.sdk.service.catalog.SchemaInfo`.
 
     Keyword Arguments:
     * `catalog_name` (str): The name of the catalog where the schema will be created. Default is `hive_metastore`.
     * `name` (str): The name of the schema. Default is a random string.
+    * `location` (str): The path to the location if it should be a managed schema.
 
     Usage:
     ```python
@@ -263,11 +261,17 @@ def test_catalog_fixture(make_catalog, make_schema, make_table):
     ```
     """
 
-    def create(*, catalog_name: str = "hive_metastore", name: str | None = None) -> SchemaInfo:
+    def create(
+        *, catalog_name: str = "hive_metastore", name: str | None = None, location: str | None = None
+    ) -> SchemaInfo:
         name = name or f"dummy_s{make_random(4)}".lower()
         full_name = f"{catalog_name}.{name}".lower()
-        sql_backend.execute(f"CREATE SCHEMA {full_name} WITH DBPROPERTIES (RemoveAfter={watchdog_remove_after})")
-        schema_info = SchemaInfo(catalog_name=catalog_name, name=name, full_name=full_name)
+        schema_ddl = f"CREATE SCHEMA {full_name}"
+        if location:
+            schema_ddl = f"{schema_ddl} LOCATION '{location}'"
+        schema_ddl = f"{schema_ddl} WITH DBPROPERTIES (RemoveAfter={watchdog_remove_after})"
+        sql_backend.execute(schema_ddl)
+        schema_info = SchemaInfo(catalog_name=catalog_name, name=name, full_name=full_name, storage_location=location)
         path = f'explore/data/{schema_info.catalog_name}/{schema_info.name}'
         log_workspace_link(f'{schema_info.full_name} schema', path)
         return schema_info

tests/integration/fixtures/test_catalog.py

@@ -17,6 +17,12 @@ def test_schema_fixture(make_schema):
     logger.info(f"Created new schema: {make_schema()}")
 
 
+def test_managed_schema_fixture(make_schema, make_random, env_or_skip):
+    schema_name = f"dummy_s{make_random(4)}".lower()
+    schema_location = f"{env_or_skip('TEST_MOUNT_CONTAINER')}/a/{schema_name}"
+    logger.info(f"Created new schema: {make_schema(location = schema_location)}")
+
+
 def test_new_managed_table_in_new_schema(make_table):
     logger.info(f"Created new managed table in new schema: {make_table()}")
 

tests/unit/fixtures/test_catalog.py

@@ -1,9 +1,15 @@
 from unittest.mock import ANY
 
-from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo
+from databricks.sdk.service.catalog import TableInfo, TableType, DataSourceFormat, FunctionInfo, SchemaInfo
 
 from databricks.labs.pytester.fixtures.unwrap import call_stateful
-from databricks.labs.pytester.fixtures.catalog import make_table, make_udf, make_catalog, make_storage_credential
+from databricks.labs.pytester.fixtures.catalog import (
+    make_table,
+    make_udf,
+    make_catalog,
+    make_storage_credential,
+    make_schema,
+)
 
 
 def test_make_table_no_args():
@@ -137,3 +143,17 @@ def test_storage_credential():
     ctx, fn_info = call_stateful(make_storage_credential, credential_name='abc')
     assert ctx is not None
     assert fn_info is not None
+
+
+def test_make_schema() -> None:
+    ctx, info = call_stateful(make_schema, name='abc', location='abfss://[email protected]')
+    assert ctx['sql_backend'].queries == [
+        "CREATE SCHEMA hive_metastore.abc LOCATION 'abfss://[email protected]' WITH DBPROPERTIES (RemoveAfter=2024091313)",
+        "DROP SCHEMA IF EXISTS hive_metastore.abc CASCADE",
+    ]
+    assert info == SchemaInfo(
+        catalog_name='hive_metastore',
+        name='abc',
+        full_name='hive_metastore.abc',
+        storage_location='abfss://[email protected]',
+    )