cleveragents-core/tests/features/steps/graph_templates_steps.py

"""
Step definitions for graph templates coverage tests.
"""

import copy
from unittest.mock import Mock, patch

from behave import given, then, when

from cleveragents.templates.base import (
    InstantiationContext,
    TemplateType,
)
from cleveragents.templates.graph_templates import GraphTemplate


@given("I have a clean test environment for graph templates")
def step_clean_environment_graph_templates(context):
    """Clean test environment setup for graph templates."""
    context.graph_template = None
    context.registry = None
    context.params = {}
    context.instantiation_context = None
    context.result = None
    context.original_definition = None


@given("I have a graph template registry")
def step_graph_template_registry(context):
    """Setup graph template registry."""
    context.registry = Mock()


@given("I have a basic graph template")
def step_basic_graph_template(context):
    """Setup basic graph template."""
    definition = {
        "name": "test_graph",
        "parameters": {"test_param": {"type": "string", "default": "test_value"}},
        "nodes": {"node1": {"type": "task", "config": {"task": "test_task"}}},
        "edges": [{"from": "node1", "to": "node2"}],
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)
    context.original_definition = copy.deepcopy(definition)


@given("I have a graph template with required parameters")
def step_graph_template_with_required_parameters(context):
    """Setup graph template with required parameters."""
    definition = {
        "name": "test_graph",
        "parameters": {
            "required_param": {"type": "string", "required": True},
            "optional_param": {"type": "string", "default": "default_value"},
        },
        "nodes": {"node1": {"type": "task", "config": {"task": "{{required_param}}"}}},
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with template variables")
def step_graph_template_with_template_variables(context):
    """Setup graph template with template variables."""
    definition = {
        "name": "{{graph_name}}",
        "nodes": {"{{node_name}}": {"type": "task", "config": {"task": "{{task_name}}"}}},
        "edges": [{"from": "{{source_node}}", "to": "{{target_node}}"}],
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template for deep copy test")
def step_graph_template_for_deep_copy_test(context):
    """Setup graph template for deep copy test."""
    definition = {
        "name": "test_graph",
        "parameters": {"test_param": {"type": "string", "default": "test_value"}},
        "nodes": {"node1": {"type": "task", "config": {"task": "test_task"}}},
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)
    context.original_definition = copy.deepcopy(definition)


@given("I have a graph template without name")
def step_graph_template_without_name(context):
    """Setup graph template without name."""
    definition = {
        "parameters": {"test_param": {"type": "string", "default": "test_value"}},
        "nodes": {"node1": {"type": "task", "config": {"task": "test_task"}}},
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with nodes")
def step_graph_template_with_nodes(context):
    """Setup graph template with various node types."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "task_node": {"type": "task", "config": {"task": "test_task"}},
            "agent_node": {"type": "agent", "agent": "test_agent"},
            "regular_node": {"type": "regular", "config": {"value": "test"}},
        },
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with None nodes")
def step_graph_template_with_none_nodes(context):
    """Setup graph template with None nodes for conditional exclusion."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "valid_node": {"type": "task", "config": {"task": "test_task"}},
            "conditional_node": None,
            "another_valid_node": {"type": "task", "config": {"task": "test_task2"}},
        },
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with agent nodes using parameter references")
def step_graph_template_with_agent_nodes_parameter_references(context):
    """Setup graph template with agent nodes using parameter references."""
    definition = {
        "name": "test_graph",
        "parameters": {"agent_ref": {"type": "string", "default": "resolved_agent"}},
        "nodes": {"agent_node": {"type": "agent", "agent": "agent_ref"}},
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with agent nodes using component references")
def step_graph_template_with_agent_nodes_component_references(context):
    """Setup graph template with agent nodes using component references."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "agent_node": {"type": "agent", "agent": "local_agent"},
            "other_node": {"type": "task", "config": {"task": "test"}},
        },
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with agent nodes using invalid references")
def step_graph_template_with_agent_nodes_invalid_references(context):
    """Setup graph template with agent nodes using invalid references."""
    definition = {
        "name": "test_graph",
        "nodes": {"agent_node": {"type": "agent", "agent": "invalid_agent_ref"}},
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with agent nodes using template variables")
def step_graph_template_with_agent_nodes_template_variables(context):
    """Setup graph template with agent nodes using template variables."""
    definition = {
        "name": "test_graph",
        "nodes": {"agent_node": {"type": "agent", "agent": "{{agent_name}}"}},
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with edges")
def step_graph_template_with_edges(context):
    """Setup graph template with various edge types."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "node1": {"type": "task", "config": {"task": "test_task1"}},
            "node2": {"type": "task", "config": {"task": "test_task2"}},
        },
        "edges": [
            {"from": "node1", "to": "node2"},
            {"from": "node2", "to": "node1", "condition": {"type": "test"}},
        ],
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with None edges")
def step_graph_template_with_none_edges(context):
    """Setup graph template with None edges for conditional exclusion."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "node1": {"type": "task", "config": {"task": "test_task1"}},
            "node2": {"type": "task", "config": {"task": "test_task2"}},
        },
        "edges": [
            {"from": "node1", "to": "node2"},
            None,
            {"from": "node2", "to": "node1"},
        ],
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with edges containing conditions")
def step_graph_template_with_edges_containing_conditions(context):
    """Setup graph template with edges containing conditions with template variables."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "node1": {"type": "task", "config": {"task": "test_task1"}},
            "node2": {"type": "task", "config": {"task": "test_task2"}},
        },
        "edges": [
            {
                "from": "node1",
                "to": "node2",
                "condition": {
                    "type": "{{condition_type}}",
                    "value": "{{condition_value}}",
                },
            }
        ],
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a graph template with edges without conditions")
def step_graph_template_with_edges_without_conditions(context):
    """Setup graph template with edges without conditions."""
    definition = {
        "name": "test_graph",
        "nodes": {
            "node1": {"type": "task", "config": {"task": "test_task1"}},
            "node2": {"type": "task", "config": {"task": "test_task2"}},
        },
        "edges": [{"from": "node1", "to": "node2"}, {"from": "node2", "to": "node1"}],
    }
    context.graph_template = GraphTemplate("test_graph", TemplateType.GRAPH, definition)


@given("I have a complete graph template with nodes and edges")
def step_complete_graph_template_with_nodes_and_edges(context):
    """Setup complete graph template with nodes and edges."""
    definition = {
        "name": "complete_graph",
        "parameters": {
            "param1": {"type": "string", "default": "value1"},
            "param2": {"type": "string", "default": "value2"},
        },
        "nodes": {
            "start_node": {"type": "task", "config": {"task": "{{param1}}"}},
            "agent_node": {"type": "agent", "agent": "test_agent"},
            "end_node": {"type": "task", "config": {"task": "{{param2}}"}},
        },
        "edges": [
            {"from": "start_node", "to": "agent_node"},
            {"from": "agent_node", "to": "end_node", "condition": {"type": "success"}},
        ],
    }
    context.graph_template = GraphTemplate("complete_graph", TemplateType.GRAPH, definition)


@given("I have an empty graph template")
def step_empty_graph_template(context):
    """Setup empty graph template."""
    definition = {"name": "empty_graph"}
    context.graph_template = GraphTemplate("empty_graph", TemplateType.GRAPH, definition)


@given("I have template parameters for graph templates")
def step_template_parameters_for_graph_templates(context):
    """Setup template parameters for graph templates."""
    context.params = {"test_param": "test_value", "optional_param": "optional_value"}


@given("I have valid template parameters for graph templates")
def step_valid_template_parameters_for_graph_templates(context):
    """Setup valid template parameters for graph templates."""
    context.params = {
        "required_param": "required_value",
        "optional_param": "optional_value",
    }


@given("I have template parameters with variable values for graph templates")
def step_template_parameters_with_variable_values_for_graph_templates(context):
    """Setup template parameters with variable values for graph templates."""
    context.params = {
        "graph_name": "dynamic_graph",
        "node_name": "dynamic_node",
        "task_name": "dynamic_task",
        "source_node": "start",
        "target_node": "end",
        "condition_type": "equals",
        "condition_value": "success",
    }


@given("I have template parameters with agent references")
def step_template_parameters_with_agent_references(context):
    """Setup template parameters with agent references."""
    context.params = {"agent_ref": "resolved_agent_name"}


@given("I have an instantiation context for graph templates")
def step_instantiation_context_for_graph_templates(context):
    """Setup instantiation context for graph templates."""
    context.instantiation_context = Mock(spec=InstantiationContext)
    context.instantiation_context.resolve_reference.return_value = None


@given("I have an instantiation context with registered agents for graph templates")
def step_instantiation_context_with_registered_agents_for_graph_templates(context):
    """Setup instantiation context with registered agents for graph templates."""
    context.instantiation_context = Mock(spec=InstantiationContext)
    agent_config = {"type": "llm", "config": {"model": "test-model"}}
    context.instantiation_context.resolve_reference.return_value = agent_config


@when("I instantiate the graph template")
def step_instantiate_graph_template(context):
    """Instantiate the graph template."""
    with patch("cleveragents.templates.graph_templates.logger") as mock_logger:
        context.mock_logger = mock_logger
        # Mock the validate_params method to return the input params
        with patch.object(context.graph_template, "validate_params", return_value=context.params):
            # Mock the _apply_template_vars method to return input unchanged
            with patch.object(
                context.graph_template,
                "_apply_template_vars",
                side_effect=lambda x, y: x,
            ):
                context.result = context.graph_template.instantiate(
                    context.params, context.registry, context.instantiation_context
                )


@then("the graph configuration should be returned")
def step_graph_configuration_should_be_returned(context):
    """Verify graph configuration is returned."""
    assert context.result is not None
    assert isinstance(context.result, dict)


@then("the parameters section should be removed from graph config")
def step_parameters_section_should_be_removed_from_graph_config(context):
    """Verify parameters section is removed from graph config."""
    assert "parameters" not in context.result


@then("template variables should be applied to the graph config")
def step_template_variables_should_be_applied_to_graph_config(context):
    """Verify template variables are applied to the graph config."""
    # This would be tested by checking if template variables are replaced
    # The actual implementation depends on the _apply_template_vars method
    assert context.result is not None


@then("the result should have correct graph structure")
def step_result_should_have_correct_graph_structure(context):
    """Verify the result has correct graph structure."""
    assert "name" in context.result
    assert context.result["name"] == "test_graph"


@then("parameter validation should be called for graph templates")
def step_parameter_validation_should_be_called_for_graph_templates(context):
    """Verify parameter validation is called for graph templates."""
    # This is tested implicitly as validate_params is called in instantiate
    assert context.result is not None


@then("the filled parameters should be used for graph templates")
def step_filled_parameters_should_be_used_for_graph_templates(context):
    """Verify filled parameters are used for graph templates."""
    # This is tested by checking that the instantiation succeeded
    assert context.result is not None


@then("template variables should be replaced with parameter values in graph config")
def step_template_variables_should_be_replaced_with_parameter_values_in_graph_config(
    context,
):
    """Verify template variables are replaced with parameter values in graph config."""
    # Check that template variables have been replaced
    assert context.result is not None
    # The actual verification would depend on the specific template variables used


@then("the original definition should not be modified for graph templates")
def step_original_definition_should_not_be_modified_for_graph_templates(context):
    """Verify the original definition is not modified for graph templates."""
    # Compare with the stored original definition
    assert context.graph_template.definition == context.original_definition


@then("the returned config should be a separate copy for graph templates")
def step_returned_config_should_be_separate_copy_for_graph_templates(context):
    """Verify the returned config is a separate copy for graph templates."""
    assert context.result is not context.graph_template.definition
    # Modify the result and ensure original is unchanged
    if "test_modification" not in context.result:
        context.result["test_modification"] = "test"
        assert "test_modification" not in context.graph_template.definition


@then("the graph name should be assigned from template name")
def step_graph_name_should_be_assigned_from_template_name(context):
    """Verify the graph name is assigned from template name."""
    assert context.result["name"] == "test_graph"


@then("nodes should be processed correctly")
def step_nodes_should_be_processed_correctly(context):
    """Verify nodes are processed correctly."""
    assert "nodes" in context.result
    assert len(context.result["nodes"]) > 0


@then("non-agent nodes should pass through unchanged")
def step_non_agent_nodes_should_pass_through_unchanged(context):
    """Verify non-agent nodes pass through unchanged."""
    nodes = context.result["nodes"]
    assert "task_node" in nodes
    assert nodes["task_node"]["type"] == "task"
    assert "regular_node" in nodes
    assert nodes["regular_node"]["type"] == "regular"


@then("None nodes should be filtered out")
def step_none_nodes_should_be_filtered_out(context):
    """Verify None nodes are filtered out."""
    nodes = context.result["nodes"]
    assert "conditional_node" not in nodes
    assert "valid_node" in nodes
    assert "another_valid_node" in nodes


@then("agent parameter references should be resolved")
def step_agent_parameter_references_should_be_resolved(context):
    """Verify agent parameter references are resolved."""
    nodes = context.result["nodes"]
    agent_node = nodes["agent_node"]
    assert agent_node["agent"] == "resolved_agent_name"


@then("agent nodes should use resolved parameter values")
def step_agent_nodes_should_use_resolved_parameter_values(context):
    """Verify agent nodes use resolved parameter values."""
    nodes = context.result["nodes"]
    agent_node = nodes["agent_node"]
    assert agent_node["agent"] == "resolved_agent_name"


@then("agent component references should be resolved")
def step_agent_component_references_should_be_resolved(context):
    """Verify agent component references are resolved."""
    # Verify that resolve_reference was called
    context.instantiation_context.resolve_reference.assert_called()


@then("agent_config should be stored for resolved agents")
def step_agent_config_should_be_stored_for_resolved_agents(context):
    """Verify agent_config is stored for resolved agents."""
    nodes = context.result["nodes"]
    agent_node = nodes["agent_node"]
    assert "agent_config" in agent_node


@then("debug logging should be called for resolved references")
def step_debug_logging_should_be_called_for_resolved_references(context):
    """Verify debug logging is called for resolved references."""
    context.mock_logger.debug.assert_called()


@then("invalid agent references should be handled gracefully")
def step_invalid_agent_references_should_be_handled_gracefully(context):
    """Verify invalid agent references are handled gracefully."""
    # The test should not fail and should return a valid result
    assert context.result is not None
    assert "nodes" in context.result


@then("agent nodes should retain original reference")
def step_agent_nodes_should_retain_original_reference(context):
    """Verify agent nodes retain original reference when resolution fails."""
    nodes = context.result["nodes"]
    agent_node = nodes["agent_node"]
    assert agent_node["agent"] == "invalid_agent_ref"


@then("template variable agent references should be preserved")
def step_template_variable_agent_references_should_be_preserved(context):
    """Verify template variable agent references are preserved."""
    # Template variables starting with {{ should not be resolved as component references
    assert context.result is not None


@then("template variables should not be resolved as components")
def step_template_variables_should_not_be_resolved_as_components(context):
    """Verify template variables are not resolved as components."""
    # The resolve_reference should not be called for template variables
    if context.instantiation_context.resolve_reference.called:
        # Check that it wasn't called with a template variable
        call_args = context.instantiation_context.resolve_reference.call_args
        if call_args:
            component_ref = call_args[0][0]
            assert not component_ref.ref_name.startswith("{{")


@then("edges should be processed correctly")
def step_edges_should_be_processed_correctly(context):
    """Verify edges are processed correctly."""
    assert "edges" in context.result
    assert len(context.result["edges"]) > 0


@then("all valid edges should be included")
def step_all_valid_edges_should_be_included(context):
    """Verify all valid edges are included."""
    edges = context.result["edges"]
    assert len(edges) == 2  # Based on the test setup


@then("None edges should be filtered out")
def step_none_edges_should_be_filtered_out(context):
    """Verify None edges are filtered out."""
    edges = context.result["edges"]
    # Should have 2 valid edges, None edge should be filtered out
    assert len(edges) == 2
    assert all(edge is not None for edge in edges)


@then("edge conditions should be processed with template variables")
def step_edge_conditions_should_be_processed_with_template_variables(context):
    """Verify edge conditions are processed with template variables."""
    edges = context.result["edges"]
    # Find the edge with condition
    edge_with_condition = next((edge for edge in edges if "condition" in edge), None)
    assert edge_with_condition is not None
    assert "condition" in edge_with_condition


@then("template variables in conditions should be replaced")
def step_template_variables_in_conditions_should_be_replaced(context):
    """Verify template variables in conditions are replaced."""
    edges = context.result["edges"]
    edge_with_condition = next((edge for edge in edges if "condition" in edge), None)
    assert edge_with_condition is not None
    # The specific replacement would depend on the _apply_template_vars implementation


@then("edges without conditions should pass through unchanged")
def step_edges_without_conditions_should_pass_through_unchanged(context):
    """Verify edges without conditions pass through unchanged."""
    edges = context.result["edges"]
    assert len(edges) == 2
    for edge in edges:
        assert "from" in edge
        assert "to" in edge


@then("the complete graph should be properly instantiated")
def step_complete_graph_should_be_properly_instantiated(context):
    """Verify the complete graph is properly instantiated."""
    assert context.result is not None
    assert "name" in context.result
    assert "nodes" in context.result
    assert "edges" in context.result


@then("all nodes should be processed")
def step_all_nodes_should_be_processed(context):
    """Verify all nodes are processed."""
    nodes = context.result["nodes"]
    assert len(nodes) == 3  # start_node, agent_node, end_node


@then("all edges should be processed")
def step_all_edges_should_be_processed(context):
    """Verify all edges are processed."""
    edges = context.result["edges"]
    assert len(edges) == 2


@then("the graph should have proper structure")
def step_graph_should_have_proper_structure(context):
    """Verify the graph has proper structure."""
    assert "name" in context.result
    assert context.result["name"] == "complete_graph"


@then("the empty graph should be handled correctly")
def step_empty_graph_should_be_handled_correctly(context):
    """Verify the empty graph is handled correctly."""
    assert context.result is not None
    assert "name" in context.result


@then("basic graph structure should be maintained")
def step_basic_graph_structure_should_be_maintained(context):
    """Verify basic graph structure is maintained."""
    assert context.result["name"] == "empty_graph"