Use ordered steps to compute step number

databrickslabs · Oct 29, 2024 · 477292a · 477292a
1 parent 1806ca3
commit 477292a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/src/databricks/labs/ucx/assessment/sequencing.py b/src/databricks/labs/ucx/assessment/sequencing.py
@@ -237,17 +237,15 @@ def generate_steps(self) -> Iterable[MigrationStep]:
         queue = self._create_node_queue(incoming)
         seen = set[MigrationNode]()
         node = queue.get()
-        step_number = 1
         ordered_steps: list[MigrationStep] = []
         while node is not None:
-            step = node.as_step(step_number, sorted(n.node_id for n in incoming[node.key]))
+            step = node.as_step(len(ordered_steps), sorted(n.node_id for n in incoming[node.key]))
             ordered_steps.append(step)
             seen.add(node)
             # Update the queue priority as if the migration step was completed
             for dependency in self._outgoing[node.key]:
                 priority = len(incoming[dependency.key] - seen)
                 queue.update(priority, dependency)
-            step_number += 1
             node = queue.get()
         return ordered_steps
 

diff --git a/tests/unit/assessment/test_sequencing.py b/tests/unit/assessment/test_sequencing.py
@@ -46,7 +46,7 @@ def get_cluster(cluster_id: str) -> ClusterDetails:
     assert steps == [
         MigrationStep(
             step_id=2,
-            step_number=1,
+            step_number=0,
             object_type="TASK",
             object_id="1234/test-task",
             object_name="test-task",
@@ -55,7 +55,7 @@ def get_cluster(cluster_id: str) -> ClusterDetails:
         ),
         MigrationStep(
             step_id=1,
-            step_number=2,
+            step_number=1,
             object_type="JOB",
             object_id="1234",
             object_name="test-job",
@@ -64,7 +64,7 @@ def get_cluster(cluster_id: str) -> ClusterDetails:
         ),
         MigrationStep(
             step_id=3,
-            step_number=3,
+            step_number=2,
             object_type="CLUSTER",
             object_id="cluster-123",
             object_name="my-cluster",