Clean up buildkite_step.py: fix import, remove dead code

buildkite/pipeline_generator/buildkite_step.py

@@ -1,6 +1,7 @@
 from pydantic import BaseModel
 from typing import Dict, List, Optional, Any, Union
 import os
+import re
 
 from step import Step
 from utils_lib.docker_utils import get_image, get_ecr_cache_registry
@@ -56,7 +57,7 @@ class BuildkiteGroupStep(BaseModel):
 def _get_step_plugin(step: Step):
     # Use K8s plugin
     use_cpu = step.device in (DeviceType.CPU, DeviceType.CPU_SMALL, DeviceType.CPU_MEDIUM)
-    if step.device in [DeviceType.H100.value, DeviceType.A100.value]:
+    if step.device in (DeviceType.H100, DeviceType.A100):
         return get_k8s_plugin(step, get_image(use_cpu))
     else:
         return {"docker#v5.2.0": get_docker_plugin(step, get_image(use_cpu))}
@@ -153,7 +154,6 @@ def _prepare_commands(step: Step, variables_to_inject: Dict[str, str]) -> List[s
             if not value:
                 continue
             # Use regex to only replace whole variable matches (not substrings)
-            import re
             # Escape variable (may have $ or special characters)
             pattern = re.escape(variable)
             command = re.sub(pattern + r'\b', value, command)
@@ -167,23 +167,19 @@ def _prepare_commands(step: Step, variables_to_inject: Dict[str, str]) -> List[s
     return final_commands
 
 
-def _create_block_step(step: Step, list_file_diff: List[str]) -> BuildkiteBlockStep:
-    block_step = BuildkiteBlockStep(
+def _create_block_step(step: Step) -> BuildkiteBlockStep:
+    return BuildkiteBlockStep(
         block=f"Run {step.label}",
         depends_on=[],
         key=f"block-{_generate_step_key(step.label)}",
     )
-    if step.label.startswith(":docker:"):
-        block_step.depends_on = []
-    return block_step
 
 
 def convert_group_step_to_buildkite_step(
     group_steps: Dict[str, List[Step]],
 ) -> List[BuildkiteGroupStep]:
     buildkite_group_steps = []
     variables_to_inject = _get_variables_to_inject()
-    print(variables_to_inject)
     global_config = get_global_config()
     list_file_diff = global_config["list_file_diff"]
 
@@ -195,7 +191,7 @@ def convert_group_step_to_buildkite_step(
             # block step
             block_step = None
             if not _step_should_run(step, list_file_diff):
-                block_step = _create_block_step(step, list_file_diff)
+                block_step = _create_block_step(step)
             if block_step:
                 group_steps_list.append(block_step)
 
@@ -235,16 +231,14 @@ def convert_group_step_to_buildkite_step(
 
             # Create AMD mirror step and its block step if specified/applicable
             if step.mirror and step.mirror.get("amd"):
-                amd_block_step = None
                 amd_block_step = BuildkiteBlockStep(
                     block=f"Run AMD: {step.label}",
                     depends_on=["image-build-amd"],
                     key=f"block-amd-{_generate_step_key(step.label)}",
                 )
                 amd_mirror_steps.append(amd_block_step)
                 amd_step = _create_amd_mirror_step(step, step_commands, step.mirror["amd"])
-                if amd_block_step:
-                    amd_step.depends_on.extend([amd_block_step.key])
+                amd_step.depends_on.extend([amd_block_step.key])
                 amd_mirror_steps.append(amd_step)
 
         buildkite_group_steps.append(

buildkite/pipeline_generator/tests/test_buildkite_step.py

@@ -0,0 +1,93 @@
+import sys, os
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+import pytest
+from unittest.mock import MagicMock
+from buildkite_step import (
+    _create_block_step,
+    _generate_step_key,
+    BuildkiteBlockStep,
+)
+from step import Step
+
+
+def _make_step(label: str, **kwargs) -> Step:
+    """Helper to create a minimal Step for testing."""
+    defaults = {
+        "label": label,
+    }
+    defaults.update(kwargs)
+    return Step(**defaults)
+
+
+# ---------- _create_block_step tests ----------
+
+class TestCreateBlockStep:
+    def test_returns_buildkite_block_step(self):
+        step = _make_step("My Test Step")
+        result = _create_block_step(step)
+        assert isinstance(result, BuildkiteBlockStep)
+
+    def test_block_field_contains_label(self):
+        step = _make_step("Lint Check")
+        result = _create_block_step(step)
+        assert result.block == "Run Lint Check"
+
+    def test_depends_on_is_empty_list(self):
+        step = _make_step("Unit Tests")
+        result = _create_block_step(step)
+        assert result.depends_on == []
+
+    def test_key_uses_generated_step_key(self):
+        step = _make_step("Build Image")
+        result = _create_block_step(step)
+        assert result.key == f"block-{_generate_step_key('Build Image')}"
+        assert result.key == "block-build-image"
+
+    def test_docker_label(self):
+        step = _make_step(":docker: Build")
+        result = _create_block_step(step)
+        assert result.block == "Run :docker: Build"
+        assert result.depends_on == []
+        assert result.key == "block--docker--build"
+
+
+# ---------- _generate_step_key tests ----------
+
+class TestGenerateStepKey:
+    def test_spaces_become_dashes(self):
+        assert _generate_step_key("hello world") == "hello-world"
+
+    def test_lowercased(self):
+        assert _generate_step_key("Hello World") == "hello-world"
+
+    def test_parentheses_removed(self):
+        assert _generate_step_key("test (gpu)") == "test-gpu"
+
+    def test_percent_removed(self):
+        assert _generate_step_key("coverage 80%") == "coverage-80"
+
+    def test_comma_becomes_dash(self):
+        assert _generate_step_key("a,b,c") == "a-b-c"
+
+    def test_plus_becomes_dash(self):
+        assert _generate_step_key("c++") == "c--"
+
+    def test_colon_becomes_dash(self):
+        assert _generate_step_key(":docker: build") == "-docker--build"
+
+    def test_dot_becomes_dash(self):
+        assert _generate_step_key("v1.2.3") == "v1-2-3"
+
+    def test_slash_becomes_dash(self):
+        assert _generate_step_key("path/to/thing") == "path-to-thing"
+
+    def test_combined_special_chars(self):
+        result = _generate_step_key("Test (GPU, 2x): v1.0+build")
+        assert result == "test-gpu--2x--v1-0-build"
+
+    def test_empty_string(self):
+        assert _generate_step_key("") == ""
+
+    def test_no_special_chars(self):
+        assert _generate_step_key("simple") == "simple"