From 246c3de7307c30f5e19d748de624904f8fbdb100 Mon Sep 17 00:00:00 2001
From: Serge Smertin <259697+nfx@users.noreply.github.com>
Date: Wed, 18 Sep 2024 13:24:02 +0200
Subject: [PATCH] [internal] port over existing UCX integration tests (#44)

---
 tests/integration/fixtures/test_catalog.py | 23 ++++++++++++
 tests/integration/fixtures/test_compute.py | 41 ++++++++++++++++++++++
 2 files changed, 64 insertions(+)

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