1  # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
  2  #
  3  # SPDX-License-Identifier: Apache-2.0
  4  
  5  import pytest
  6  
  7  from haystack.components.agents import Agent
  8  from haystack.core.component.component import component
  9  from haystack.tools import Tool, Toolset
 10  from haystack.tools.toolset import _ToolsetWrapper
 11  
 12  
 13  # Test fixtures
 14  def add_numbers(a: int, b: int) -> int:
 15      """Add two numbers."""
 16      return a + b
 17  
 18  
 19  def multiply_numbers(a: int, b: int) -> int:
 20      """Multiply two numbers."""
 21      return a * b
 22  
 23  
 24  def subtract_numbers(a: int, b: int) -> int:
 25      """Subtract b from a."""
 26      return a - b
 27  
 28  
 29  @pytest.fixture
 30  def add_tool():
 31      return Tool(
 32          name="add",
 33          description="Add two numbers",
 34          parameters={
 35              "type": "object",
 36              "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
 37              "required": ["a", "b"],
 38          },
 39          function=add_numbers,
 40      )
 41  
 42  
 43  @pytest.fixture
 44  def multiply_tool():
 45      return Tool(
 46          name="multiply",
 47          description="Multiply two numbers",
 48          parameters={
 49              "type": "object",
 50              "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
 51              "required": ["a", "b"],
 52          },
 53          function=multiply_numbers,
 54      )
 55  
 56  
 57  @pytest.fixture
 58  def subtract_tool():
 59      return Tool(
 60          name="subtract",
 61          description="Subtract two numbers",
 62          parameters={
 63              "type": "object",
 64              "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
 65              "required": ["a", "b"],
 66          },
 67          function=subtract_numbers,
 68      )
 69  
 70  
 71  class TestToolsetWrapper:
 72      """Tests for the _ToolsetWrapper class"""
 73  
 74      def test_toolset_plus_toolset_creates_wrapper(self, add_tool, multiply_tool):
 75          """Test that combining two Toolsets creates a _ToolsetWrapper and works correctly."""
 76          toolset1 = Toolset([add_tool])
 77          toolset2 = Toolset([multiply_tool])
 78  
 79          result = toolset1 + toolset2
 80  
 81          assert isinstance(result, _ToolsetWrapper)
 82          assert len(result) == 2
 83          assert add_tool in result
 84          assert multiply_tool in result
 85  
 86      def test_wrapper_with_agent(self, add_tool, multiply_tool):
 87          """Test that _ToolsetWrapper works with Agent."""
 88  
 89          @component
 90          class MockChatGenerator:
 91              def run(self, messages, tools=None, **kwargs):
 92                  return {"replies": messages}
 93  
 94              def warm_up(self):
 95                  pass
 96  
 97          toolset1 = Toolset([add_tool])
 98          toolset2 = Toolset([multiply_tool])
 99          wrapper = toolset1 + toolset2
100  
101          agent = Agent(chat_generator=MockChatGenerator(), tools=wrapper)
102          agent.warm_up()
103  
104          assert len(list(agent.tools)) == 2
105  
106      def test_wrapper_chaining_and_duplicate_detection(self, add_tool, multiply_tool, subtract_tool):
107          """Test chaining operations and that duplicates are still detected."""
108          toolset1 = Toolset([add_tool])
109          toolset2 = Toolset([multiply_tool])
110          toolset3 = Toolset([subtract_tool])
111  
112          # Chaining should work
113          result = toolset1 + toolset2 + toolset3
114          assert len(result) == 3
115  
116          # Duplicates should be detected
117          toolset_with_dup = Toolset([add_tool])
118          with pytest.raises(ValueError, match="Duplicate tool names found"):
119              _ = result + toolset_with_dup