diff --git a/tests/integration/fixtures/test_catalog.py b/tests/integration/fixtures/test_catalog.py index d6b3939..a028a62 100644 --- a/tests/integration/fixtures/test_catalog.py +++ b/tests/integration/fixtures/test_catalog.py @@ -2,6 +2,9 @@ import pytest + + + logger = logging.getLogger(__name__) @@ -56,3 +59,23 @@ def test_storage_credential(env_or_skip, make_storage_credential, make_random): credential_name=credential_name, aws_iam_role_arn=env_or_skip("TEST_UBER_ROLE_ID"), ) + + +def test_remove_after_property_table(ws, make_table, sql_backend): + new_table = make_table() + # TODO: tables.get is currently failing with + # databricks.sdk.errors.platform.NotFound: Catalog 'hive_metastore' does not exist. + sql_response = list(sql_backend.fetch(f"DESCRIBE TABLE EXTENDED {new_table.full_name}")) + for row in sql_response: + if row.col_name == "Table Properties": + assert "RemoveAfter" in row[1] + + +def test_remove_after_property_schema(ws, make_schema, sql_backend): + new_schema = make_schema() + # TODO: schemas.get is currently failing with + # databricks.sdk.errors.platform.NotFound: Catalog 'hive_metastore' does not exist. + sql_response = list(sql_backend.fetch(f"DESCRIBE SCHEMA EXTENDED {new_schema.full_name}")) + for row in sql_response: + if row.database_description_item == "Properties": + assert "RemoveAfter" in row[1] diff --git a/tests/integration/fixtures/test_compute.py b/tests/integration/fixtures/test_compute.py index 41d53fd..141d9ea 100644 --- a/tests/integration/fixtures/test_compute.py +++ b/tests/integration/fixtures/test_compute.py @@ -1,7 +1,10 @@ import logging +from datetime import datetime, timedelta, timezone from databricks.sdk.service.iam import PermissionLevel +from databricks.labs.pytester.fixtures.watchdog import TEST_RESOURCE_PURGE_TIMEOUT + logger = logging.getLogger(__name__) @@ -36,3 +39,41 @@ def test_warehouse_has_remove_after_tag(ws, make_warehouse): created_warehouse = ws.warehouses.get(new_warehouse.response.id) warehouse_tags = created_warehouse.tags.as_dict() assert warehouse_tags["custom_tags"][0]["key"] == "RemoveAfter" + + +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) + job_tags = created_job.settings.tags + assert "RemoveAfter" in job_tags + + purge_time = datetime.strptime(job_tags["RemoveAfter"], "%Y%m%d%H").replace(tzinfo=timezone.utc) + assert (purge_time - datetime.now(timezone.utc)) < (TEST_RESOURCE_PURGE_TIMEOUT + timedelta(hours=1)) # noqa: F405 + + +def test_remove_after_tag_clusters(ws, env_or_skip, make_cluster): + new_cluster = make_cluster(single_node=True, instance_pool_id=env_or_skip('TEST_INSTANCE_POOL_ID')) + created_cluster = ws.clusters.get(new_cluster.cluster_id) + cluster_tags = created_cluster.custom_tags + assert "RemoveAfter" in cluster_tags + purge_time = datetime.strptime(cluster_tags["RemoveAfter"], "%Y%m%d%H").replace(tzinfo=timezone.utc) + assert (purge_time - datetime.now(timezone.utc)) < (TEST_RESOURCE_PURGE_TIMEOUT + timedelta(hours=1)) # noqa: F405 + + +def test_remove_after_tag_warehouse(ws, env_or_skip, make_warehouse): + new_warehouse = make_warehouse() + created_warehouse = ws.warehouses.get(new_warehouse.response.id) + warehouse_tags = created_warehouse.tags.as_dict() + assert warehouse_tags["custom_tags"][0]["key"] == "RemoveAfter" + remove_after_tag = warehouse_tags["custom_tags"][0]["value"] + purge_time = datetime.strptime(remove_after_tag, "%Y%m%d%H").replace(tzinfo=timezone.utc) + assert (purge_time - datetime.now(timezone.utc)) < (TEST_RESOURCE_PURGE_TIMEOUT + timedelta(hours=1)) # noqa: F405 + + +def test_remove_after_tag_instance_pool(ws, make_instance_pool): + new_instance_pool = make_instance_pool() + created_instance_pool = ws.instance_pools.get(new_instance_pool.instance_pool_id) + pool_tags = created_instance_pool.custom_tags + assert "RemoveAfter" in pool_tags + purge_time = datetime.strptime(pool_tags["RemoveAfter"], "%Y%m%d%H").replace(tzinfo=timezone.utc) + assert (purge_time - datetime.now(timezone.utc)) < (TEST_RESOURCE_PURGE_TIMEOUT + timedelta(hours=1)) # noqa: F405