yamlgraph 0.3.9__py3-none-any.whl

tests/unit/test_shell_tools.py

@@ -0,0 +1,229 @@
+"""Tests for shell tool executor."""
+
+from yamlgraph.tools.shell import (
+    ShellToolConfig,
+    execute_shell_tool,
+    parse_tools,
+    sanitize_variables,
+)
+
+
+class TestShellToolConfig:
+    """Tests for ShellToolConfig dataclass."""
+
+    def test_default_values(self):
+        """Config has sensible defaults."""
+        config = ShellToolConfig(command="echo hello")
+        assert config.command == "echo hello"
+        assert config.description == ""
+        assert config.parse == "text"
+        assert config.timeout == 30
+        assert config.working_dir == "."
+        assert config.env == {}
+        assert config.success_codes == [0]
+
+    def test_custom_values(self):
+        """Config accepts custom values."""
+        config = ShellToolConfig(
+            command="curl http://api.example.com",
+            description="Fetch API data",
+            parse="json",
+            timeout=60,
+            working_dir="/tmp",
+            env={"API_KEY": "secret"},
+            success_codes=[0, 1],
+        )
+        assert config.parse == "json"
+        assert config.timeout == 60
+        assert config.env == {"API_KEY": "secret"}
+
+
+class TestSanitizeVariables:
+    """Tests for variable sanitization."""
+
+    def test_sanitizes_simple_string(self):
+        """Simple strings are quoted."""
+        result = sanitize_variables({"name": "Alice"})
+        # shlex.quote adds quotes around strings with no special chars
+        assert result["name"] in ("Alice", "'Alice'")
+
+    def test_sanitizes_shell_injection(self):
+        """Shell injection attempts are safely quoted."""
+        # Command substitution attempt
+        result = sanitize_variables({"name": "$(rm -rf /)"})
+        assert "$" not in result["name"] or result["name"].startswith("'")
+        # The result should be a quoted string
+        assert result["name"] == "'$(rm -rf /)'"
+
+    def test_sanitizes_semicolon_injection(self):
+        """Semicolon command chaining is prevented."""
+        result = sanitize_variables({"name": "test; rm -rf /"})
+        assert "'" in result["name"]  # Must be quoted
+
+    def test_sanitizes_pipe_injection(self):
+        """Pipe injection is prevented."""
+        result = sanitize_variables({"name": "test | cat /etc/passwd"})
+        assert "'" in result["name"]  # Must be quoted
+
+    def test_handles_none_values(self):
+        """None values become empty strings."""
+        result = sanitize_variables({"name": None})
+        assert result["name"] == ""
+
+    def test_handles_list_values(self):
+        """List values are JSON encoded and quoted."""
+        result = sanitize_variables({"items": [1, 2, 3]})
+        assert "[1, 2, 3]" in result["items"] or result["items"] == "'[1, 2, 3]'"
+
+
+class TestExecuteShellTool:
+    """Tests for execute_shell_tool function."""
+
+    def test_executes_command(self):
+        """Simple command executes successfully."""
+        config = ShellToolConfig(command="echo hello")
+        result = execute_shell_tool(config, {})
+        assert result.success is True
+        assert result.output.strip() == "hello"
+        assert result.error is None
+
+    def test_substitutes_variables(self):
+        """Placeholders replaced with values."""
+        config = ShellToolConfig(command="echo {message}")
+        result = execute_shell_tool(config, {"message": "world"})
+        assert result.success is True
+        assert result.output.strip() == "world"
+
+    def test_multiple_variables(self):
+        """Multiple placeholders all substituted."""
+        config = ShellToolConfig(command="echo {a} {b} {c}")
+        result = execute_shell_tool(config, {"a": "1", "b": "2", "c": "3"})
+        assert result.output.strip() == "1 2 3"
+
+    def test_parses_json_output(self):
+        """JSON stdout parsed to dict."""
+        # Double braces escape them from .format()
+        config = ShellToolConfig(
+            command='echo \'{{"name": "test", "value": 42}}\'',
+            parse="json",
+        )
+        result = execute_shell_tool(config, {})
+        assert result.success is True
+        assert result.output == {"name": "test", "value": 42}
+
+    def test_parse_none_returns_none(self):
+        """parse=none returns None for side-effect commands."""
+        config = ShellToolConfig(command="echo ignored", parse="none")
+        result = execute_shell_tool(config, {})
+        assert result.success is True
+        assert result.output is None
+
+    def test_handles_timeout(self):
+        """Long-running command times out."""
+        config = ShellToolConfig(command="sleep 10", timeout=1)
+        result = execute_shell_tool(config, {})
+        assert result.success is False
+        assert "timed out" in result.error.lower()
+
+    def test_captures_stderr_on_error(self):
+        """Non-zero exit captures stderr."""
+        config = ShellToolConfig(command="ls /nonexistent_path_xyz")
+        result = execute_shell_tool(config, {})
+        assert result.success is False
+        assert result.error is not None
+        assert "No such file" in result.error or "nonexistent" in result.error.lower()
+
+    def test_custom_success_codes(self):
+        """Custom success codes treated as success."""
+        # grep returns 1 when no match found
+        config = ShellToolConfig(
+            command="grep nonexistent /dev/null",
+            success_codes=[0, 1],
+        )
+        result = execute_shell_tool(config, {})
+        assert result.success is True
+
+    def test_working_dir(self):
+        """Command runs in specified directory."""
+        config = ShellToolConfig(command="pwd", working_dir="/tmp")
+        result = execute_shell_tool(config, {})
+        assert result.success is True
+        assert "/tmp" in result.output
+
+    def test_env_variables(self):
+        """Environment variables passed to command."""
+        config = ShellToolConfig(
+            command="echo $TEST_VAR",
+            env={"TEST_VAR": "secret_value"},
+        )
+        result = execute_shell_tool(config, {})
+        assert result.success is True
+        assert "secret_value" in result.output
+
+    def test_invalid_json_parse_fails(self):
+        """Invalid JSON returns error."""
+        config = ShellToolConfig(command="echo 'not json'", parse="json")
+        result = execute_shell_tool(config, {})
+        assert result.success is False
+        assert "json" in result.error.lower()
+
+
+class TestParseTools:
+    """Tests for parse_tools function."""
+
+    def test_empty_config(self):
+        """Empty config returns empty registry."""
+        registry = parse_tools({})
+        assert registry == {}
+
+    def test_parses_single_tool(self):
+        """Single tool parsed correctly."""
+        config = {
+            "search": {
+                "command": "curl -s {url}",
+                "description": "Search the web",
+                "parse": "json",
+                "timeout": 60,
+            }
+        }
+        registry = parse_tools(config)
+        assert "search" in registry
+        assert registry["search"].command == "curl -s {url}"
+        assert registry["search"].description == "Search the web"
+        assert registry["search"].parse == "json"
+        assert registry["search"].timeout == 60
+
+    def test_parses_multiple_tools(self):
+        """Multiple tools all parsed."""
+        config = {
+            "tool1": {"command": "echo 1"},
+            "tool2": {"command": "echo 2"},
+            "tool3": {"command": "echo 3"},
+        }
+        registry = parse_tools(config)
+        assert len(registry) == 3
+        assert all(name in registry for name in ["tool1", "tool2", "tool3"])
+
+    def test_default_values_applied(self):
+        """Missing optional fields get defaults."""
+        config = {"minimal": {"command": "echo hello"}}
+        registry = parse_tools(config)
+        tool = registry["minimal"]
+        assert tool.description == ""
+        assert tool.parse == "text"
+        assert tool.timeout == 30
+        assert tool.working_dir == "."
+        assert tool.env == {}
+
+    def test_parses_env_and_working_dir(self):
+        """env and working_dir parsed correctly."""
+        config = {
+            "script": {
+                "command": "node index.js",
+                "working_dir": "./scripts",
+                "env": {"NODE_ENV": "production"},
+            }
+        }
+        registry = parse_tools(config)
+        assert registry["script"].working_dir == "./scripts"
+        assert registry["script"].env == {"NODE_ENV": "production"}

tests/unit/test_state_builder.py

@@ -0,0 +1,331 @@
+"""Unit tests for dynamic state builder.
+
+TDD: Red phase - these tests define the expected behavior.
+"""
+
+from operator import add
+from typing import Annotated, get_args, get_origin
+
+from yamlgraph.models.state_builder import sorted_add
+
+
+class TestSortedAdd:
+    """Test the sorted_add reducer for map node fan-in."""
+
+    def test_concatenates_lists(self):
+        """Should concatenate two lists."""
+        result = sorted_add([1, 2], [3, 4])
+        assert result == [1, 2, 3, 4]
+
+    def test_handles_empty_existing(self):
+        """Should handle empty existing list."""
+        result = sorted_add([], [1, 2])
+        assert result == [1, 2]
+
+    def test_handles_none_existing(self):
+        """Should handle None as existing list."""
+        result = sorted_add(None, [1, 2])
+        assert result == [1, 2]
+
+    def test_handles_empty_new(self):
+        """Should handle empty new list."""
+        result = sorted_add([1, 2], [])
+        assert result == [1, 2]
+
+    def test_handles_none_new(self):
+        """Should handle None as new list."""
+        result = sorted_add([1, 2], None)
+        assert result == [1, 2]
+
+    def test_sorts_by_map_index(self):
+        """Should sort results by _map_index for map fan-in."""
+        # Simulate out-of-order parallel results
+        existing = [{"_map_index": 2, "value": "third"}]
+        new = [{"_map_index": 0, "value": "first"}]
+        result = sorted_add(existing, new)
+
+        assert result[0]["_map_index"] == 0
+        assert result[0]["value"] == "first"
+        assert result[1]["_map_index"] == 2
+        assert result[1]["value"] == "third"
+
+    def test_sorts_multiple_out_of_order(self):
+        """Should sort many out-of-order items correctly."""
+        # Simulate 5 items arriving in random order
+        items = [
+            {"_map_index": 3, "data": "d"},
+            {"_map_index": 0, "data": "a"},
+            {"_map_index": 4, "data": "e"},
+            {"_map_index": 1, "data": "b"},
+            {"_map_index": 2, "data": "c"},
+        ]
+        result = sorted_add([], items)
+
+        assert [r["data"] for r in result] == ["a", "b", "c", "d", "e"]
+
+    def test_no_sort_for_non_dict_items(self):
+        """Should not sort if items are not dicts."""
+        result = sorted_add([3, 1], [2])
+        assert result == [3, 1, 2]  # Preserved insertion order
+
+    def test_no_sort_for_dicts_without_map_index(self):
+        """Should not sort if dicts lack _map_index."""
+        result = sorted_add([{"a": 1}], [{"b": 2}])
+        assert result == [{"a": 1}, {"b": 2}]
+
+
+class TestBuildStateClass:
+    """Test dynamic TypedDict generation from graph config."""
+
+    def test_includes_base_infrastructure_fields(self):
+        """State always has infrastructure fields."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        annotations = State.__annotations__
+        assert "thread_id" in annotations
+        assert "current_step" in annotations
+        assert "errors" in annotations
+        assert "messages" in annotations
+
+    def test_errors_has_reducer(self):
+        """errors field uses Annotated[list, add] reducer."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        errors_type = State.__annotations__["errors"]
+        assert get_origin(errors_type) is Annotated
+        args = get_args(errors_type)
+        assert args[0] is list
+        assert args[1] is add
+
+    def test_messages_has_reducer(self):
+        """messages field uses Annotated[list, add] reducer."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        messages_type = State.__annotations__["messages"]
+        assert get_origin(messages_type) is Annotated
+        args = get_args(messages_type)
+        assert args[0] is list
+        assert args[1] is add
+
+    def test_extracts_state_key_from_nodes(self):
+        """state_key in node config becomes state field."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {
+            "nodes": {
+                "generate": {"prompt": "generate", "state_key": "generated"},
+                "analyze": {"prompt": "analyze", "state_key": "analysis"},
+            },
+            "edges": [],
+        }
+        State = build_state_class(config)
+
+        assert "generated" in State.__annotations__
+        assert "analysis" in State.__annotations__
+
+    def test_agent_node_adds_input_field(self):
+        """Agent nodes automatically add 'input' field."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {
+            "nodes": {
+                "agent": {"type": "agent", "prompt": "agent"},
+            },
+            "edges": [],
+        }
+        State = build_state_class(config)
+
+        assert "input" in State.__annotations__
+
+    def test_agent_node_adds_tool_results_field(self):
+        """Agent nodes add _tool_results field."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {
+            "nodes": {
+                "agent": {"type": "agent", "prompt": "agent"},
+            },
+            "edges": [],
+        }
+        State = build_state_class(config)
+
+        assert "_tool_results" in State.__annotations__
+
+    def test_router_node_adds_route_field(self):
+        """Router nodes add _route field."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {
+            "nodes": {
+                "router": {
+                    "type": "router",
+                    "prompt": "router",
+                    "routes": {"a": "node_a", "b": "node_b"},
+                },
+            },
+            "edges": [],
+        }
+        State = build_state_class(config)
+
+        assert "_route" in State.__annotations__
+
+    def test_loop_tracking_fields_included(self):
+        """Loop tracking fields are always included."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        assert "_loop_counts" in State.__annotations__
+        assert "_loop_limit_reached" in State.__annotations__
+        assert "_agent_iterations" in State.__annotations__
+        assert "_agent_limit_reached" in State.__annotations__
+
+    def test_state_is_typeddict_total_false(self):
+        """Generated state is TypedDict with total=False (all optional)."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        # TypedDict with total=False has __total__ = False
+        assert State.__total__ is False
+
+    def test_state_works_with_langgraph(self):
+        """Generated state class works with LangGraph StateGraph."""
+        from langgraph.graph import StateGraph
+
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {
+            "nodes": {
+                "test": {"prompt": "test", "state_key": "result"},
+            },
+            "edges": [],
+        }
+        State = build_state_class(config)
+
+        # Should not raise
+        graph = StateGraph(State)
+        graph.add_node("test", lambda s: {"result": "done"})
+        graph.set_entry_point("test")
+        graph.set_finish_point("test")
+        compiled = graph.compile()
+
+        # Verify fields are preserved
+        result = compiled.invoke({"input": "hello"})
+        assert "result" in result
+
+    def test_reducer_accumulates_messages(self):
+        """Messages reducer accumulates across nodes."""
+        from langgraph.graph import StateGraph
+
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        graph = StateGraph(State)
+        graph.add_node("n1", lambda s: {"messages": [{"content": "a"}]})
+        graph.add_node("n2", lambda s: {"messages": [{"content": "b"}]})
+        graph.add_edge("n1", "n2")
+        graph.set_entry_point("n1")
+        graph.set_finish_point("n2")
+        compiled = graph.compile()
+
+        result = compiled.invoke({})
+        assert len(result["messages"]) == 2
+
+
+class TestExtractNodeFields:
+    """Test field extraction from node configurations."""
+
+    def test_extracts_state_key(self):
+        """Extracts state_key from nodes."""
+        from yamlgraph.models.state_builder import extract_node_fields
+
+        nodes = {
+            "gen": {"state_key": "generated"},
+            "analyze": {"state_key": "analysis"},
+        }
+        fields = extract_node_fields(nodes)
+
+        assert "generated" in fields
+        assert "analysis" in fields
+
+    def test_agent_adds_special_fields(self):
+        """Agent nodes add input and _tool_results."""
+        from yamlgraph.models.state_builder import extract_node_fields
+
+        nodes = {"agent": {"type": "agent"}}
+        fields = extract_node_fields(nodes)
+
+        assert "input" in fields
+        assert "_tool_results" in fields
+
+    def test_router_adds_route_field(self):
+        """Router nodes add _route."""
+        from yamlgraph.models.state_builder import extract_node_fields
+
+        nodes = {"router": {"type": "router", "routes": {}}}
+        fields = extract_node_fields(nodes)
+
+        assert "_route" in fields
+
+
+class TestCommonInputFields:
+    """Test that common input fields are included."""
+
+    def test_includes_topic_field(self):
+        """topic field included for content generation."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        assert "topic" in State.__annotations__
+
+    def test_includes_style_field(self):
+        """style field included for content generation."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        assert "style" in State.__annotations__
+
+    def test_includes_word_count_field(self):
+        """word_count field included for content generation."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        assert "word_count" in State.__annotations__
+
+    def test_includes_message_field(self):
+        """message field included for router."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        assert "message" in State.__annotations__
+
+    def test_includes_input_field(self):
+        """input field included for agents."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        config = {"nodes": {}, "edges": []}
+        State = build_state_class(config)
+
+        assert "input" in State.__annotations__

tests/unit/test_state_builder_map.py

@@ -0,0 +1,104 @@
+"""Tests for state_builder map node reducer support."""
+
+from typing import Annotated, get_args, get_origin
+
+from yamlgraph.models.state_builder import (
+    build_state_class,
+    extract_node_fields,
+    sorted_add,
+)
+
+
+class TestExtractNodeFieldsMap:
+    """Tests for map node collect field extraction."""
+
+    def test_map_node_collect_field_added(self) -> None:
+        """Map node adds collect field to extracted fields."""
+        nodes = {
+            "expand_frames": {
+                "type": "map",
+                "over": "{state.frames}",
+                "sub_node": "expand_frame",
+                "collect": "expanded_frames",
+            }
+        }
+        fields = extract_node_fields(nodes)
+        assert "expanded_frames" in fields
+
+    def test_map_node_collect_has_sorted_reducer(self) -> None:
+        """Map node collect field has Annotated[list, sorted_add] type."""
+        nodes = {
+            "expand_frames": {
+                "type": "map",
+                "over": "{state.frames}",
+                "sub_node": "expand_frame",
+                "collect": "expanded_frames",
+            }
+        }
+        fields = extract_node_fields(nodes)
+
+        field_type = fields["expanded_frames"]
+
+        # Check it's Annotated
+        assert get_origin(field_type) is Annotated
+        # Check args: list and sorted_add for ordered fan-in
+        args = get_args(field_type)
+        assert args[0] is list
+        assert args[1] is sorted_add
+
+    def test_map_node_without_collect_no_field(self) -> None:
+        """Map node without collect key doesn't add field."""
+        nodes = {
+            "expand_frames": {
+                "type": "map",
+                "over": "{state.frames}",
+                "sub_node": "expand_frame",
+                # No collect key
+            }
+        }
+        fields = extract_node_fields(nodes)
+        # Should not have any fields from this node
+        # (other fields may come from sub_node if it had output_key)
+        assert "expanded_frames" not in fields
+
+
+class TestBuildStateClassMap:
+    """Tests for build_state_class with map nodes."""
+
+    def test_build_state_includes_collect_field(self) -> None:
+        """Built state class includes map node collect field."""
+        config = {
+            "nodes": {
+                "expand_frames": {
+                    "type": "map",
+                    "over": "{state.frames}",
+                    "sub_node": "expand_frame",
+                    "collect": "expanded_frames",
+                }
+            }
+        }
+        state_class = build_state_class(config)
+        annotations = state_class.__annotations__
+
+        assert "expanded_frames" in annotations
+
+    def test_build_state_collect_has_sorted_reducer(self) -> None:
+        """Built state class has sorted_add reducer for collect field."""
+        config = {
+            "nodes": {
+                "expand_frames": {
+                    "type": "map",
+                    "over": "{state.frames}",
+                    "sub_node": "expand_frame",
+                    "collect": "expanded_frames",
+                }
+            }
+        }
+        state_class = build_state_class(config)
+        annotations = state_class.__annotations__
+
+        field_type = annotations["expanded_frames"]
+        assert get_origin(field_type) is Annotated
+        args = get_args(field_type)
+        assert args[0] is list
+        assert args[1] is sorted_add