diff --git a/buildkite/pipeline_generator/buildkite_step.py b/buildkite/pipeline_generator/buildkite_step.py index ff31d99d..7f7a8610 100644 --- a/buildkite/pipeline_generator/buildkite_step.py +++ b/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,15 +167,12 @@ 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( @@ -183,7 +180,6 @@ def convert_group_step_to_buildkite_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,7 +231,6 @@ 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"], @@ -243,8 +238,7 @@ def convert_group_step_to_buildkite_step( ) 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( diff --git a/buildkite/pipeline_generator/tests/__init__.py b/buildkite/pipeline_generator/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/buildkite/pipeline_generator/tests/test_buildkite_step.py b/buildkite/pipeline_generator/tests/test_buildkite_step.py new file mode 100644 index 00000000..82f03ffb --- /dev/null +++ b/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"