test: reorganize WebSocket tool tests with comprehensive coverage

Co-Authored-By: Joe Moura <joao@crewai.com>

tests/tools/test_tool_usage.py

@@ -146,60 +146,159 @@ class MockWebSocketTool(BaseTool):
         return self._run(**input)
 
 
-def test_websocket_tool_repeated_usage():
+class TestWebSocketToolUsage:
-    tool = MockWebSocketTool()
+    """Test cases for WebSocket tool usage and duplicate detection."""
-    agent = Agent(
-        role="Test Agent",
-        goal="Test WebSocket tools",
-        backstory="Testing WebSocket tool execution",
-        tools=[tool],
-        verbose=True,
-    )
 
-    task = Task(
+    @pytest.fixture
-        description="Test WebSocket tool",
+    def setup_websocket_tool(self):
-        expected_output="Test output",
+        """Fixture to set up WebSocket tool and agent for testing."""
-        agent=agent,
+        tool = MockWebSocketTool()
-    )
+        agent = Agent(
+            role="Test Agent",
+            goal="Test WebSocket tools",
+            backstory="Testing WebSocket tool execution",
+            tools=[tool],
+            verbose=True,
+        )
+        return tool, agent
 
-    tool_usage = ToolUsage(
+    def test_first_execution(self, setup_websocket_tool):
-        tools_handler=ToolsHandler(),
+        """Test first execution of WebSocket tool."""
-        tools=[tool],
+        tool, agent = setup_websocket_tool
-        original_tools=[tool],
+        task = Task(
-        tools_description="WebSocket tool for testing",
+            description="Test WebSocket tool",
-        tools_names="websocket_tool",
+            expected_output="Test output",
-        task=task,
+            agent=agent,
-        function_calling_llm=MagicMock(),
+        )
-        agent=agent,
+        tool_usage = ToolUsage(
-        action=MagicMock(),
+            tools_handler=ToolsHandler(),
-    )
+            tools=[tool],
+            original_tools=[tool],
+            tools_description="WebSocket tool for testing",
+            tools_names="websocket_tool",
+            task=task,
+            function_calling_llm=MagicMock(),
+            agent=agent,
+            action=MagicMock(),
+        )
+        calling = ToolCalling(
+            tool_name="WebSocket Tool",
+            arguments={"question": "Test question"},
+            log="Test log",
+        )
+        result = tool_usage.use(calling, "Test string")
+        assert "Answer to: Test question" in result
 
-    # First call
+    def test_repeated_execution(self, setup_websocket_tool):
-    calling1 = ToolCalling(
+        """Test repeated execution with same question is detected."""
-        tool_name="WebSocket Tool",
+        tool, agent = setup_websocket_tool
-        arguments={"question": "Test question"},
+        task = Task(
-        log="Test log",
+            description="Test WebSocket tool",
-    )
+            expected_output="Test output",
-    result1 = tool_usage.use(calling1, "Test string")
+            agent=agent,
-    assert "Answer to: Test question" in result1
+        )
+        tool_usage = ToolUsage(
+            tools_handler=ToolsHandler(),
+            tools=[tool],
+            original_tools=[tool],
+            tools_description="WebSocket tool for testing",
+            tools_names="websocket_tool",
+            task=task,
+            function_calling_llm=MagicMock(),
+            agent=agent,
+            action=MagicMock(),
+        )
+        # First call
+        calling1 = ToolCalling(
+            tool_name="WebSocket Tool",
+            arguments={"question": "Test question"},
+            log="Test log",
+        )
+        result1 = tool_usage.use(calling1, "Test string")
+        assert "Answer to: Test question" in result1
 
-    # Same question should be detected as repeated
+        # Same question should be detected as repeated
-    calling2 = ToolCalling(
+        calling2 = ToolCalling(
-        tool_name="WebSocket Tool",
+            tool_name="WebSocket Tool",
-        arguments={"question": "Test question"},
+            arguments={"question": "Test question"},
-        log="Test log",
+            log="Test log",
-    )
+        )
-    result2 = tool_usage.use(calling2, "Test string")
+        result2 = tool_usage.use(calling2, "Test string")
-    assert "reusing the same input" in result2.lower()
+        assert "reusing the same input" in result2.lower()
 
-    # Different question should work
+    def test_different_question(self, setup_websocket_tool):
-    calling3 = ToolCalling(
+        """Test execution with different questions works."""
-        tool_name="WebSocket Tool",
+        tool, agent = setup_websocket_tool
-        arguments={"question": "Different question"},
+        task = Task(
-        log="Test log",
+            description="Test WebSocket tool",
-    )
+            expected_output="Test output",
-    result3 = tool_usage.use(calling3, "Test string")
+            agent=agent,
-    assert "Answer to: Different question" in result3
+        )
+        tool_usage = ToolUsage(
+            tools_handler=ToolsHandler(),
+            tools=[tool],
+            original_tools=[tool],
+            tools_description="WebSocket tool for testing",
+            tools_names="websocket_tool",
+            task=task,
+            function_calling_llm=MagicMock(),
+            agent=agent,
+            action=MagicMock(),
+        )
+        # First question
+        calling1 = ToolCalling(
+            tool_name="WebSocket Tool",
+            arguments={"question": "First question"},
+            log="Test log",
+        )
+        result1 = tool_usage.use(calling1, "Test string")
+        assert "Answer to: First question" in result1
+
+        # Different question should work
+        calling2 = ToolCalling(
+            tool_name="WebSocket Tool",
+            arguments={"question": "Second question"},
+            log="Test log",
+        )
+        result2 = tool_usage.use(calling2, "Test string")
+        assert "Answer to: Second question" in result2
+
+    def test_invalid_arguments(self, setup_websocket_tool):
+        """Test handling of invalid arguments."""
+        tool, agent = setup_websocket_tool
+        task = Task(
+            description="Test WebSocket tool",
+            expected_output="Test output",
+            agent=agent,
+        )
+        tool_usage = ToolUsage(
+            tools_handler=ToolsHandler(),
+            tools=[tool],
+            original_tools=[tool],
+            tools_description="WebSocket tool for testing",
+            tools_names="websocket_tool",
+            task=task,
+            function_calling_llm=MagicMock(),
+            agent=agent,
+            action=MagicMock(),
+        )
+        # Test with empty arguments
+        calling = ToolCalling(
+            tool_name="WebSocket Tool",
+            arguments={},
+            log="Test log",
+        )
+        with pytest.raises(ValueError):
+            tool_usage.use(calling, "Test string")
+
+        # Test with None arguments
+        calling = ToolCalling(
+            tool_name="WebSocket Tool",
+            arguments=None,
+            log="Test log",
+        )
+        result = tool_usage.use(calling, "Test string")
+        assert "error" in result.lower()
 
 
 def test_validate_tool_input_booleans_and_none():