fix(make_warehouse): accept custom tags #220

Author: b0lle

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

@@ -339,6 +339,7 @@ def create(
         cluster_size: str | None = None,
         max_num_clusters: int = 1,
         enable_serverless_compute: bool = False,
+        tags: list[EndpointTagPair] = [],
         **kwargs,
     ) -> Wait[GetWarehouseResponse]:
         if warehouse_name is None:
@@ -348,7 +349,9 @@ def create(
         if cluster_size is None:
             cluster_size = "2X-Small"
 
-        remove_after_tags = EndpointTags(custom_tags=[EndpointTagPair(key="RemoveAfter", value=watchdog_remove_after)])
+        remove_after_tags = EndpointTags(
+            custom_tags=[EndpointTagPair(key="RemoveAfter", value=watchdog_remove_after)] + tags
+        )
         return ws.warehouses.create(
             name=warehouse_name,
             cluster_size=cluster_size,

tests/integration/fixtures/test_compute.py

@@ -4,6 +4,7 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.iam import PermissionLevel
 from databricks.sdk.service.jobs import RunResultState, SparkPythonTask
+from databricks.sdk.service.sql import EndpointTagPair
 
 from databricks.labs.pytester.fixtures.watchdog import TEST_RESOURCE_PURGE_TIMEOUT
 
@@ -55,6 +56,13 @@ def test_warehouse_has_remove_after_tag(ws, make_warehouse):
     assert warehouse_tags["custom_tags"][0]["key"] == "RemoveAfter"
 
 
+def test_warehouse_has_custom_tag(ws, make_warehouse):
+    new_warehouse = make_warehouse(tags=[EndpointTagPair(key="my-custom-tag", value="my-custom-tag-value")])
+    created_warehouse = ws.warehouses.get(new_warehouse.response.id)
+    warehouse_tags = created_warehouse.tags.as_dict()
+    assert warehouse_tags["custom_tags"][1]["key"] == "my-custom-tag"
+
+
 def test_remove_after_tag_jobs(ws, env_or_skip, make_job):
     new_job = make_job()
     created_job = ws.jobs.get(new_job.job_id)