Improving resiliency of integration tests (#594)

Author: nfx

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
 test = [
     "coverage[toml]>=6.5",
     "pytest",
+    "pytest-xdist",
     "pytest-cov>=4.0.0,<5.0.0",
     "pytest-mock>=3.0.0,<4.0.0",
 ]
@@ -65,6 +66,7 @@ dependencies = [
 
 [tool.hatch.envs.integration.scripts]
 test = "pytest --cov src tests/integration"
+parallel = "pytest -n auto --cov src tests/integration"
 
 [tool.hatch.envs.lint]
 detached = true

src/databricks/labs/ucx/mixins/sql.py

@@ -5,12 +5,15 @@
 from collections.abc import Iterator
 from datetime import timedelta
 
+from databricks.sdk import errors
+from databricks.sdk.errors import NotFound
 from databricks.sdk.service.sql import (
     ColumnInfoTypeName,
     Disposition,
     ExecuteStatementResponse,
     Format,
     ResultData,
+    ServiceErrorCode,
     StatementExecutionAPI,
     StatementState,
     StatementStatus,
@@ -91,10 +94,28 @@ def __init__(self, api_client):
     def _raise_if_needed(status: StatementStatus):
         if status.state not in [StatementState.FAILED, StatementState.CANCELED, StatementState.CLOSED]:
             return
-        msg = status.state.value
-        if status.error is not None:
-            msg = f"{msg}: {status.error.error_code.value} {status.error.message}"
-        raise RuntimeError(msg)
+        if "SCHEMA_NOT_FOUND" in status.error.message:
+            raise NotFound(status.error.message)
+        if "TABLE_NOT_FOUND" in status.error.message:
+            raise NotFound(status.error.message)
+        mapping = {
+            ServiceErrorCode.ABORTED: errors.Aborted,
+            ServiceErrorCode.ALREADY_EXISTS: errors.AlreadyExists,
+            ServiceErrorCode.BAD_REQUEST: errors.BadRequest,
+            ServiceErrorCode.CANCELLED: errors.Cancelled,
+            ServiceErrorCode.DEADLINE_EXCEEDED: errors.DeadlineExceeded,
+            ServiceErrorCode.INTERNAL_ERROR: errors.InternalError,
+            ServiceErrorCode.IO_ERROR: errors.InternalError,
+            ServiceErrorCode.NOT_FOUND: errors.NotFound,
+            ServiceErrorCode.RESOURCE_EXHAUSTED: errors.ResourceExhausted,
+            ServiceErrorCode.SERVICE_UNDER_MAINTENANCE: errors.TemporarilyUnavailable,
+            ServiceErrorCode.TEMPORARILY_UNAVAILABLE: errors.TemporarilyUnavailable,
+            ServiceErrorCode.UNAUTHENTICATED: errors.Unauthenticated,
+            ServiceErrorCode.UNKNOWN: errors.Unknown,
+            ServiceErrorCode.WORKSPACE_TEMPORARILY_UNAVAILABLE: errors.TemporarilyUnavailable,
+        }
+        error_class = mapping.get(status.error.error_code, errors.Unknown)
+        raise error_class(status.error.message)
 
     def execute(
         self,

src/databricks/labs/ucx/workspace_access/groups.py

@@ -312,6 +312,7 @@ def _get_or_create_backup_group(self, source_group_name: str, source_group: iam.
             roles=source_group.roles,
             members=source_group.members,
         )  # TODO: there still could be a corner case, where we get `Group with name db-temp-XXX already exists.`
+        # TODO: (nfx) except ResourceConflict: ... find group
         self._workspace_groups.append(backup_group)
         logger.info(f"Backup group {backup_group_name} successfully created")
 

tests/integration/assessment/test_assessment.py

@@ -1,5 +1,8 @@
 import logging
+from datetime import timedelta
 
+from databricks.sdk.errors import NotFound
+from databricks.sdk.retries import retried
 from databricks.sdk.service import compute, jobs
 
 from databricks.labs.ucx.assessment.crawlers import (
@@ -121,6 +124,7 @@ def test_spn_crawler(ws, inventory_schema, make_job, make_pipeline, sql_backend)
     assert results[0].tenant_id == _TEST_TENANT_ID
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=5))
 def test_spn_crawler_no_config(ws, inventory_schema, make_job, make_pipeline, sql_backend, make_cluster):
     make_job()
     make_pipeline()
@@ -167,6 +171,7 @@ def test_spn_crawler_deleted_cluster_policy(
     assert any(_ for _ in results if _.storage_account == _TEST_STORAGE_ACCOUNT)
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=5))
 def test_spn_crawler_with_pipeline_unavailable_secret(ws, inventory_schema, make_job, make_pipeline, sql_backend):
     make_job(spark_conf=_SPARK_CONF)
     make_pipeline(configuration=_PIPELINE_CONF_WITH_SECRET)

tests/integration/framework/test_crawlers.py

@@ -8,4 +8,4 @@ def test_deploys_database(sql_backend, inventory_schema):
     deployer = SchemaDeployer(sql_backend, inventory_schema, ucx)
     deployer.deploy_schema()
     deployer.deploy_table("grants", Grant)
-    deployer.deploy_view("grant_detail", "assessment/views/grant_detail.sql")
+    deployer.deploy_view("grant_detail", "queries/views/grant_detail.sql")

tests/integration/framework/test_fixtures.py

@@ -1,7 +1,10 @@
 import logging
+from datetime import timedelta
 
 import pytest
 from _pytest.outcomes import Failed, Skipped
+from databricks.sdk.errors import NotFound
+from databricks.sdk.retries import retried
 from databricks.sdk.service.workspace import AclPermission
 
 from databricks.labs.ucx.mixins.compute import CommandExecutor
@@ -73,6 +76,7 @@ def test_pipeline(make_pipeline):
     logger.info(f"created {make_pipeline()}")
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=5))
 def test_this_wheel_installs(ws, wsfs_wheel):
     commands = CommandExecutor(ws)
 

tests/integration/hive_metastore/test_grants.py

@@ -1,10 +1,15 @@
 import logging
+from datetime import timedelta
+
+from databricks.sdk.errors import NotFound
+from databricks.sdk.retries import retried
 
 from databricks.labs.ucx.hive_metastore import GrantsCrawler, TablesCrawler
 
 logger = logging.getLogger(__name__)
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=10))
 def test_all_grants_in_databases(sql_backend, inventory_schema, make_schema, make_table, make_group):
     group_a = make_group()
     group_b = make_group()
@@ -36,6 +41,7 @@ def test_all_grants_in_databases(sql_backend, inventory_schema, make_schema, mak
         all_grants[f"{grant.principal}.{grant.object_key}"] = grant.action_type
 
     assert len(all_grants) >= 8, "must have at least three grants"
+    # TODO: (nfx) KeyError: 'sdk-xH7i.hive_metastore.default'
     assert all_grants[f"{group_a.display_name}.hive_metastore.default"] == "USAGE"
     assert all_grants[f"{group_b.display_name}.hive_metastore.default"] == "USAGE"
     assert all_grants[f"{group_a.display_name}.{table_a.full_name}"] == "SELECT"

tests/integration/test_installation.py

@@ -18,6 +18,7 @@
 logger = logging.getLogger(__name__)
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=10))
 def test_job_failure_propagates_correct_error_message_and_logs(ws, sql_backend, env_or_skip, make_random):
     default_cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID")
     tacl_cluster_id = env_or_skip("TEST_LEGACY_TABLE_ACL_CLUSTER_ID")

tests/integration/workspace_access/test_groups.py

@@ -53,6 +53,7 @@ def test_prepare_environment(ws, make_ucx_group):
             continue
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=5))
 def test_prepare_environment_no_groups_selected(ws, make_ucx_group, make_group, make_acc_group):
     make_group()
     make_acc_group()
@@ -117,6 +118,8 @@ def test_replace_workspace_groups_with_account_groups(
     permission_manager.inventorize_permissions()
 
     dummy_grants = list(permission_manager.load_all_for("TABLE", dummy_table.full_name, Grant))
+    # TODO: (nfx) where 1 = len([Grant(principal='ucx_BjI1', action_type='MODIFY', catalog='hive_metastore',
+    #  database='ucx_s4ygj', table='ucx_tstdr', view=None, any_file=False, anonymous_function=False)])
     assert 2 == len(dummy_grants)
 
     table_permissions = grants.for_table_info(dummy_table)
@@ -285,6 +288,7 @@ def test_migration_state_should_be_saved_without_missing_anything(sql_backend, m
     assert len(new_state.groups) == len(state.groups)
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=10))
 def test_set_owner_permission(
     ws,
     sql_backend,
@@ -316,6 +320,8 @@ def test_set_owner_permission(
     permission_manager.inventorize_permissions()
 
     dummy_grants = list(permission_manager.load_all_for("TABLE", dummy_table.full_name, Grant))
+    # TODO (nfx) where 1 = len([Grant(principal='ucx_Z1Ga', action_type='MODIFY', catalog='hive_metastore',
+    #  database='ucx_syubc', table='ucx_tiinm', view=None, any_file=False, anonymous_function=False)])
     assert 2 == len(dummy_grants)
 
     table_permissions = grants.for_table_info(dummy_table)

tests/integration/workspace_access/test_scim.py

@@ -10,6 +10,7 @@
 from databricks.labs.ucx.workspace_access.scim import ScimSupport
 
 
+# TODO: (nfx) databricks.sdk.errors.mapping.ResourceConflict: None Group with name db-temp-ucx_QSFR already exists.
 def test_scim(ws: WorkspaceClient, make_ucx_group, sql_backend, inventory_schema):
     """
     This test does the following: