Improve code quality with better import sorting and test organization

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

tests/flow/test_flow_tool_result.py

@@ -12,6 +12,7 @@ from crewai.tools import BaseTool
 class TestToolInput(BaseModel):
     query: str = Field(..., description='Query to process')
 
+
 class TestTool(BaseTool):
     name: str = 'Test Tool'
     description: str = 'A test tool to demonstrate the issue'
@@ -21,63 +22,112 @@ class TestTool(BaseTool):
     def _run(self, query: str) -> str:
         return f'Result for query: {query}'
 
-def test_flow_tool_result_as_answer():
-    """Test that tools with result_as_answer=True are properly processed in Flow mode."""
-    # Create a test tool
-    test_tool = TestTool()
-    
-    # Create a test agent with the tool
-    agent = Agent(
-        role='Tester',
-        goal='Test tools',
-        backstory='Testing tools in Flow vs Crew',
-        tools=[test_tool]
-    )
 
-    # Create a task with the tool
-    task = Task(
-        description='Test task using the tool',
-        expected_output='Test output',
-        agent=agent,
-        tools=[test_tool]
-    )
+class TestFlowToolResult:
+    """Test suite for Flow tool result handling."""
     
-    # Create a simple Flow with direct access to the agent
-    class SimpleFlow(Flow):
-        def __init__(self):
-            super().__init__()
-            self.test_agent = agent
-            
-        @start()
-        def start_task(self):
-            return 'Task started'
-            
-        @listen('start_task')
-        @Flow.with_agent(agent)  # Associate the agent with this method
-        def execute_task(self):
-            # Simulate tool execution and setting tools_results
-            self.test_agent.tools_results = [{
-                "name": "Test Tool",
-                "input": {"query": "test"},
-                "result": "Result for query: test",
-                "result_as_answer": True
-            }]
-            return "Agent task execution result"
+    @pytest.fixture
+    def test_tool(self):
+        """Create a test tool with result_as_answer=True."""
+        return TestTool()
+    
+    @pytest.fixture
+    def test_agent(self, test_tool):
+        """Create a test agent with the tool."""
+        return Agent(
+            role='Tester',
+            goal='Test tools',
+            backstory='Testing tools in Flow vs Crew',
+            tools=[test_tool]
+        )
+    
+    @pytest.fixture
+    def test_task(self, test_agent, test_tool):
+        """Create a task with the tool."""
+        return Task(
+            description='Test task using the tool',
+            expected_output='Test output',
+            agent=test_agent,
+            tools=[test_tool]
+        )
+    
+    def test_flow_tool_result_as_answer(self, test_agent, test_task):
+        """Test that tools with result_as_answer=True are properly processed in Flow mode."""
+        # Create a simple Flow with direct access to the agent
+        class SimpleFlow(Flow):
+            def __init__(self):
+                super().__init__()
+                self.test_agent = test_agent
+                
+            @start()
+            def start_task(self):
+                return 'Task started'
+                
+            @listen('start_task')
+            @Flow.with_agent(test_agent)  # Associate the agent with this method
+            def execute_task(self):
+                # Simulate tool execution and setting tools_results
+                self.test_agent.tools_results = [{
+                    "name": "Test Tool",
+                    "input": {"query": "test"},
+                    "result": "Result for query: test",
+                    "result_as_answer": True
+                }]
+                return "Agent task execution result"
 
-    # Create a mock for Crew to return the same result
-    with patch('crewai.crew.Crew.kickoff') as mock_crew_kickoff:
-        mock_crew_kickoff.return_value = "Result for query: test"
-        
-        # Test Flow
+        # Create a mock for Crew to return the same result
+        with patch('crewai.crew.Crew.kickoff') as mock_crew_kickoff:
+            mock_crew_kickoff.return_value = "Result for query: test"
+            
+            # Test Flow
+            flow = SimpleFlow()
+            flow_result = flow.kickoff()
+            
+            # Verify that Flow returns the tool result with our fix
+            assert flow_result == "Result for query: test", "Flow should return tool result with result_as_answer=True"
+            
+            # Test Crew
+            crew = Crew(agents=[test_agent], tasks=[test_task])
+            crew_result = crew.kickoff()
+            
+            # Verify that Crew returns the tool result
+            assert crew_result == "Result for query: test", "Crew should return tool result with result_as_answer=True"
+    
+    def test_flow_multiple_tool_results(self, test_agent, test_task):
+        """Test handling of multiple tool results with result_as_answer=True."""
+        # Create a simple Flow with direct access to the agent
+        class SimpleFlow(Flow):
+            def __init__(self):
+                super().__init__()
+                self.test_agent = test_agent
+                
+            @start()
+            def start_task(self):
+                return 'Task started'
+                
+            @listen('start_task')
+            @Flow.with_agent(test_agent)
+            def execute_task(self):
+                # Simulate multiple tool executions with result_as_answer=True
+                self.test_agent.tools_results = [
+                    {
+                        "name": "Test Tool 1",
+                        "input": {"query": "test1"},
+                        "result": "Result for query: test1",
+                        "result_as_answer": True
+                    },
+                    {
+                        "name": "Test Tool 2",
+                        "input": {"query": "test2"},
+                        "result": "Result for query: test2",
+                        "result_as_answer": True
+                    }
+                ]
+                return "Agent task execution result"
+                
+        # Test Flow with multiple tool results
         flow = SimpleFlow()
         flow_result = flow.kickoff()
         
-        # Verify that Flow returns the tool result with our fix
-        assert flow_result == "Result for query: test", "Flow should return tool result with result_as_answer=True"
-        
-        # Test Crew
-        crew = Crew(agents=[agent], tasks=[task])
-        crew_result = crew.kickoff()
-        
-        # Verify that Crew returns the tool result
-        assert crew_result == "Result for query: test", "Crew should return tool result with result_as_answer=True"
+        # Verify that Flow returns the first tool result with result_as_answer=True
+        assert flow_result == "Result for query: test1", "Flow should return the first tool result with result_as_answer=True"