tests/test_utils.py

import pytest

from workflows.errors import CyclicDependencyError, UnknownVariableError, WorkflowError
from workflows.utils import _create_variable_step_mapping, create_dependency_graph, topological_sort


# Dummy classes to simulate Workflow, Step, and Field
class DummyField:
    def __init__(self, name, type="str", variable=None):
        self.name = name
        self.type = type
        # For input fields, variable property is needed
        self.variable = variable if variable is not None else name


class DummyStep:
    def __init__(self, input_fields, output_fields):
        self.input_fields = input_fields
        self.output_fields = output_fields


class DummyWorkflow:
    def __init__(self, steps):
        # steps is a dict with key as step_id and value as DummyStep
        self.steps = steps


# Tests for _create_variable_step_mapping


def test_create_variable_step_mapping_success():
    # Create a workflow with two steps producing unique output variables
    step_a = DummyStep(input_fields=[], output_fields=[DummyField("out1")])
    step_b = DummyStep(input_fields=[], output_fields=[DummyField("out2")])
    workflow = DummyWorkflow({"A": step_a, "B": step_b})
    mapping = _create_variable_step_mapping(workflow)
    assert mapping == {"A.out1": "A", "B.out2": "B"}


def test_create_variable_step_mapping_duplicate():
    # Create a workflow where two steps produce an output with same name
    step_a = DummyStep(input_fields=[], output_fields=[DummyField("out"), DummyField("out")])
    workflow = DummyWorkflow({"A": step_a})
    with pytest.raises(WorkflowError):
        _create_variable_step_mapping(workflow)


def test_create_variable_step_mapping_empty():
    """Test _create_variable_step_mapping with an empty workflow should return an empty mapping."""
    workflow = DummyWorkflow({})
    mapping = _create_variable_step_mapping(workflow)
    assert mapping == {}


def test_create_variable_step_mapping_multiple_outputs():
    """Test a workflow where a single step produces multiple outputs with unique names."""
    step = DummyStep(input_fields=[], output_fields=[DummyField("out1"), DummyField("out2")])
    workflow = DummyWorkflow({"A": step})
    mapping = _create_variable_step_mapping(workflow)
    assert mapping == {"A.out1": "A", "A.out2": "A"}


# Tests for create_dependency_graph


def test_create_dependency_graph_success_with_dependency():
    # Step A produces 'A.out', which is used as input in step B
    step_a = DummyStep(input_fields=[], output_fields=[DummyField("out")])
    # For input_fields, explicitly set variable to reference A.out
    step_b = DummyStep(input_fields=[DummyField("dummy", variable="A.out")], output_fields=[DummyField("result")])
    workflow = DummyWorkflow({"A": step_a, "B": step_b})
    # No external input provided for A.out so dependency must be created
    deps = create_dependency_graph(workflow, input_values={})
    # Step B depends on step A
    assert deps["B"] == {"A"}
    # Step A has no dependencies
    assert deps["A"] == set()


def test_create_dependency_graph_success_with_external_input():
    # Step B expects an input, but it is provided externally
    step_b = DummyStep(
        input_fields=[DummyField("param", variable="external_param")], output_fields=[DummyField("result")]
    )
    workflow = DummyWorkflow({"B": step_b})
    # Provide external input for external_param
    deps = create_dependency_graph(workflow, input_values={"external_param": 42})
    # With external input, no dependency is needed
    assert deps["B"] == set()


def test_create_dependency_graph_unknown_variable():
    # Step B expects an input that is neither produced by any step nor provided externally
    step_b = DummyStep(
        input_fields=[DummyField("param", variable="non_existent")], output_fields=[DummyField("result")]
    )
    workflow = DummyWorkflow({"B": step_b})
    with pytest.raises(UnknownVariableError):
        _ = create_dependency_graph(workflow, input_values={})


def test_create_dependency_graph_complex():
    """Test create_dependency_graph on a more complex workflow with multiple dependencies."""
    # Step A produces A.out, Step B uses A.out, Step C uses B.out, and Step D uses both A.out and B.out
    step_a = DummyStep(input_fields=[], output_fields=[DummyField("out")])
    step_b = DummyStep(input_fields=[DummyField("inp", variable="A.out")], output_fields=[DummyField("out")])
    step_c = DummyStep(input_fields=[DummyField("inp", variable="B.out")], output_fields=[DummyField("result")])
    step_d = DummyStep(
        input_fields=[DummyField("inp1", variable="A.out"), DummyField("inp2", variable="B.out")],
        output_fields=[DummyField("final")],
    )

    workflow = DummyWorkflow({"A": step_a, "B": step_b, "C": step_c, "D": step_d})
    # Provide external input for "B.out" so that step B's output isn't expected to come from a step
    # However, to simulate dependency, assume external input is not provided for the dependencies used in step C and D
    # Therefore, workflow must resolve A.out for step B, and then step B produces B.out for steps C and D.
    # Let's not provide any external input, so both dependencies are created.

    deps = create_dependency_graph(workflow, input_values={})
    # Expected dependencies:
    # B depends on A
    # C depends on B
    # D depends on both A and B
    assert deps["B"] == {"A"}
    assert deps["C"] == {"B"}
    assert deps["D"] == {"A", "B"}


# Tests for topological_sort


def test_topological_sort_success():
    # Create a simple dependency graph: A -> B -> C
    deps = {"A": set(), "B": {"A"}, "C": {"B"}}
    order = topological_sort(deps)
    # Check that order satisfies dependencies: A before B, B before C
    assert order.index("A") < order.index("B") < order.index("C")


def test_topological_sort_cycle():
    # Create a cyclic dependency: A -> B and B -> A
    deps = {"A": {"B"}, "B": {"A"}}
    with pytest.raises(CyclicDependencyError):
        _ = topological_sort(deps)


def test_topological_sort_single_node():
    """Test topological_sort on a graph with a single node and no dependencies."""
    deps = {"A": set()}
    order = topological_sort(deps)
    assert order == ["A"]


def test_topological_sort_disconnected():
    """Test topological_sort on a graph with disconnected nodes (no dependencies among them)."""
    deps = {"A": set(), "B": set(), "C": set()}
    order = topological_sort(deps)
    # The order can be in any permutation, but must contain all nodes
    assert set(order) == {"A", "B", "C"}