Skip to content

Commit 6cc819a

Browse files
authored
Clean up for workspace scan tables (#2788)
Had some dangling commits for #2754
1 parent cf20a40 commit 6cc819a

File tree

2 files changed

+45
-6
lines changed

2 files changed

+45
-6
lines changed

tests/integration/progress/test_workflows.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,29 @@
1-
from datetime import timedelta
1+
import datetime as dt
22

33
from databricks.sdk.errors import NotFound, InvalidParameterValue
44
from databricks.sdk.retries import retried
55

66
from ..conftest import MockInstallationContext
77

88

9-
@retried(on=[NotFound, InvalidParameterValue], timeout=timedelta(minutes=12))
9+
@retried(on=[NotFound, InvalidParameterValue], timeout=dt.timedelta(minutes=12))
1010
def test_running_real_migration_progress_job(installation_ctx: MockInstallationContext) -> None:
1111
"""Ensure that the migration-progress workflow can complete successfully."""
1212
installation_ctx.workspace_installation.run()
1313

1414
# The assessment workflow is a prerequisite for migration-progress: it needs to successfully complete before we can
1515
# test the migration-progress workflow.
16-
installation_ctx.deployed_workflows.run_workflow("assessment")
17-
assert installation_ctx.deployed_workflows.validate_step("assessment")
16+
workflow = "assessment"
17+
installation_ctx.deployed_workflows.run_workflow(workflow, max_wait=dt.timedelta(minutes=25))
18+
assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
1819

1920
# After the assessment, a user (maybe) installs the progress tracking
2021
installation_ctx.progress_tracking_installation.run()
2122

2223
# Run the migration-progress workflow until completion.
23-
installation_ctx.deployed_workflows.run_workflow("migration-progress-experimental")
24-
assert installation_ctx.deployed_workflows.validate_step("migration-progress-experimental")
24+
workflow = "migration-progress-experimental"
25+
installation_ctx.deployed_workflows.run_workflow(workflow)
26+
assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
2527

2628
# Ensure that the migration-progress workflow populated the `workflow_runs` table.
2729
query = f"SELECT 1 FROM {installation_ctx.ucx_catalog}.multiworkspace.workflow_runs"
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import datetime as dt
2+
import pytest
3+
4+
from databricks.labs.ucx.config import WorkspaceConfig
5+
from databricks.labs.ucx.contexts.workflow_task import RuntimeContext
6+
from databricks.labs.ucx.progress.workflow_runs import WorkflowRunRecorder
7+
8+
9+
class MockWorkspaceClient:
10+
"""Mock workspace client"""
11+
12+
def get_workspace_id(self) -> int:
13+
return 123456789
14+
15+
16+
@pytest.mark.parametrize(
17+
"attribute, klass",
18+
[
19+
("workflow_run_recorder", WorkflowRunRecorder),
20+
],
21+
)
22+
def test_context_attribute_class(tmp_path, attribute, klass) -> None:
23+
"""Increase coverage with these tests"""
24+
config = WorkspaceConfig("str")
25+
named_parameters = {
26+
"workflow": "test",
27+
"job_id": "123",
28+
"parent_run_id": "456",
29+
"start_time": dt.datetime.now(tz=dt.timezone.utc).replace(microsecond=0).isoformat(),
30+
}
31+
32+
ctx = RuntimeContext(named_parameters=named_parameters).replace(
33+
config=config, sql_backend=None, workspace_client=MockWorkspaceClient()
34+
)
35+
36+
assert hasattr(ctx, attribute)
37+
assert isinstance(getattr(ctx, attribute), klass)

0 commit comments

Comments
 (0)