Clean up for workspace scan tables

tests/integration/progress/test_workflows.py

@@ -1,27 +1,29 @@
-from datetime import timedelta
+import datetime as dt
 
 from databricks.sdk.errors import NotFound, InvalidParameterValue
 from databricks.sdk.retries import retried
 
 from ..conftest import MockInstallationContext
 
 
-@retried(on=[NotFound, InvalidParameterValue], timeout=timedelta(minutes=12))
+@retried(on=[NotFound, InvalidParameterValue], timeout=dt.timedelta(minutes=12))
 def test_running_real_migration_progress_job(installation_ctx: MockInstallationContext) -> None:
     """Ensure that the migration-progress workflow can complete successfully."""
     installation_ctx.workspace_installation.run()
 
     # The assessment workflow is a prerequisite for migration-progress: it needs to successfully complete before we can
     # test the migration-progress workflow.
-    installation_ctx.deployed_workflows.run_workflow("assessment")
-    assert installation_ctx.deployed_workflows.validate_step("assessment")
+    workflow = "assessment"
+    installation_ctx.deployed_workflows.run_workflow(workflow, max_wait=dt.timedelta(minutes=25))
+    assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
 
     # After the assessment, a user (maybe) installs the progress tracking
     installation_ctx.progress_tracking_installation.run()
 
     # Run the migration-progress workflow until completion.
-    installation_ctx.deployed_workflows.run_workflow("migration-progress-experimental")
-    assert installation_ctx.deployed_workflows.validate_step("migration-progress-experimental")
+    workflow = "migration-progress-experimental"
+    installation_ctx.deployed_workflows.run_workflow(workflow)
+    assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
 
     # Ensure that the migration-progress workflow populated the `workflow_runs` table.
     query = f"SELECT 1 FROM {installation_ctx.ucx_catalog}.multiworkspace.workflow_runs"

tests/unit/contexts/test_workflow_task.py

@@ -0,0 +1,37 @@
+import datetime as dt
+import pytest
+
+from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.ucx.contexts.workflow_task import RuntimeContext
+from databricks.labs.ucx.progress.workflow_runs import WorkflowRunRecorder
+
+
+class MockWorkspaceClient:
+    """Mock workspace client"""
+
+    def get_workspace_id(self) -> int:
+        return 123456789
+
+
+@pytest.mark.parametrize(
+    "attribute, klass",
+    [
+        ("workflow_run_recorder", WorkflowRunRecorder),
+    ],
+)
+def test_context_attribute_class(tmp_path, attribute, klass) -> None:
+    """Increase coverage with these tests"""
+    config = WorkspaceConfig("str")
+    named_parameters = {
+        "workflow": "test",
+        "job_id": "123",
+        "parent_run_id": "456",
+        "start_time": dt.datetime.now(tz=dt.timezone.utc).replace(microsecond=0).isoformat(),
+    }
+
+    ctx = RuntimeContext(named_parameters=named_parameters).replace(
+        config=config, sql_backend=None, workspace_client=MockWorkspaceClient()
+    )
+
+    assert hasattr(ctx, attribute)
+    assert isinstance(getattr(ctx, attribute), klass)