codeflash-agent/packages/codeflash-api/tests/test_workflow.py

from __future__ import annotations

import pytest
from pydantic import ValidationError

from codeflash_api.workflow._router import _extract_yaml_steps
from codeflash_api.workflow.schemas import (
    WorkflowGenErrorResponse,
    WorkflowGenRequest,
    WorkflowGenResponse,
)


class TestWorkflowGenRequest:
    """Tests for WorkflowGenRequest schema."""

    def test_minimal_valid(self) -> None:
        """
        Empty dicts are accepted; defaults fill in trace_id and config.
        """
        req = WorkflowGenRequest(
            repo_files={},
            directory_structure={},
        )
        assert {} == req.repo_files
        assert {} == req.directory_structure
        assert req.codeflash_config is None
        assert "" == req.trace_id

    def test_all_fields(self) -> None:
        """
        All fields are accepted when provided.
        """
        req = WorkflowGenRequest(
            repo_files={"setup.py": "from setuptools import setup"},
            directory_structure={"src": {"main.py": None}},
            codeflash_config={"language": "python"},
            trace_id="abc-123",
        )
        assert {"setup.py": "from setuptools import setup"} == req.repo_files
        assert {"src": {"main.py": None}} == req.directory_structure
        assert {"language": "python"} == req.codeflash_config
        assert "abc-123" == req.trace_id

    def test_missing_repo_files_rejected(self) -> None:
        """
        Omitting repo_files raises a validation error.
        """
        with pytest.raises(ValidationError):
            WorkflowGenRequest(
                directory_structure={},
            )  # type: ignore[call-arg]

    def test_missing_directory_structure_rejected(self) -> None:
        """
        Omitting directory_structure raises a validation error.
        """
        with pytest.raises(ValidationError):
            WorkflowGenRequest(
                repo_files={},
            )  # type: ignore[call-arg]


class TestWorkflowGenResponse:
    """Tests for WorkflowGenResponse schema."""

    def test_workflow_steps(self) -> None:
        """
        Response carries the workflow_steps string.
        """
        resp = WorkflowGenResponse(workflow_steps="steps:\n- name: Checkout")
        assert "steps:\n- name: Checkout" == resp.workflow_steps

    def test_missing_workflow_steps_rejected(self) -> None:
        """
        Omitting workflow_steps raises a validation error.
        """
        with pytest.raises(ValidationError):
            WorkflowGenResponse()  # type: ignore[call-arg]


class TestWorkflowGenErrorResponse:
    """Tests for WorkflowGenErrorResponse schema."""

    def test_error_message(self) -> None:
        """
        Error response carries the error string.
        """
        resp = WorkflowGenErrorResponse(error="bad input")
        assert "bad input" == resp.error

    def test_missing_error_rejected(self) -> None:
        """
        Omitting error raises a validation error.
        """
        with pytest.raises(ValidationError):
            WorkflowGenErrorResponse()  # type: ignore[call-arg]


class TestExtractYamlSteps:
    """Tests for _extract_yaml_steps YAML extraction logic."""

    def test_full_steps_section(self) -> None:
        """
        Extracts steps content when 'steps:' header is present.
        """
        text = (
            "steps:\n"
            "- name: Checkout\n"
            "  uses: actions/checkout@v4\n"
            "- name: Setup Python\n"
            "  uses: actions/setup-python@v5\n"
        )
        result = _extract_yaml_steps(text)
        assert result is not None
        assert result.startswith("steps:\n")
        assert "- name: Checkout" in result
        assert "- name: Setup Python" in result

    def test_steps_section_with_indented_dashes(self) -> None:
        """
        Lines with leading whitespace before '-' get normalized to 2-space indent.
        """
        text = "steps:\n    - name: Checkout\n    - name: Install\n"
        result = _extract_yaml_steps(text)
        assert result is not None
        assert "  - name: Checkout" in result
        assert "  - name: Install" in result

    def test_non_dash_lines_preserved(self) -> None:
        """
        Lines that don't start with '-' are kept as-is.
        """
        text = (
            "steps:\n"
            "- name: Checkout\n"
            "  uses: actions/checkout@v4\n"
            "  with:\n"
            "    fetch-depth: 0\n"
        )
        result = _extract_yaml_steps(text)
        assert result is not None
        assert "  uses: actions/checkout@v4" in result
        assert "  with:" in result

    def test_fallback_array_match(self) -> None:
        """
        Falls back to array_match_re when no 'steps:' header is found.
        """
        text = (
            "Here is the output:\n\n"
            "- name: Checkout\n"
            "  uses: actions/checkout@v4\n"
            "- name: Install deps\n"
            "  run: pip install -r requirements.txt\n"
        )
        result = _extract_yaml_steps(text)
        assert result is not None
        assert result.startswith("steps:\n")
        assert "- name: Checkout" in result

    def test_fallback_yaml_like(self) -> None:
        """
        Falls back to yaml_like_re as last resort.
        """
        text = "Some preamble text\n- name: Only step\n  run: echo hello"
        result = _extract_yaml_steps(text)
        assert result is not None
        assert result.startswith("steps:\n")
        assert "- name: Only step" in result

    def test_no_yaml_returns_none(self) -> None:
        """
        Returns None when no YAML-like content is found.
        """
        result = _extract_yaml_steps("No yaml here at all.")
        assert result is None

    def test_empty_string_returns_none(self) -> None:
        """
        Returns None for empty input.
        """
        result = _extract_yaml_steps("")
        assert result is None

    def test_steps_header_only_returns_bare_header(self) -> None:
        """
        Returns bare 'steps:\\n' when header has no content lines.
        """
        result = _extract_yaml_steps("steps:\n")
        assert "steps:\n" == result

    def test_blank_line_terminates_steps_capture(self) -> None:
        """
        The steps regex stops capturing at the first blank line.
        """
        text = "steps:\n- name: A\n\n- name: B\n"
        result = _extract_yaml_steps(text)
        assert result is not None
        assert "- name: A" in result
        assert "- name: B" not in result

    def test_markdown_fenced_content_with_steps(self) -> None:
        """
        Extracts steps even when surrounded by markdown fences.
        """
        text = (
            "```yaml\n"
            "steps:\n"
            "- name: Checkout\n"
            "  uses: actions/checkout@v4\n"
            "```\n"
        )
        result = _extract_yaml_steps(text)
        assert result is not None
        assert "- name: Checkout" in result

    def test_double_newline_terminates_array_match(self) -> None:
        """
        The array_match_re stops at a double newline boundary.
        """
        text = "- name: First\n  run: echo 1\n\n\nSome extra text after"
        result = _extract_yaml_steps(text)
        assert result is not None
        assert "Some extra text" not in result