From 5775ed3fcb5c51b92e6d57071a36c2697a799b0a Mon Sep 17 00:00:00 2001
From: Brandon Hancock <brandon@brandonhancock.io>
Date: Sun, 23 Jun 2024 09:42:33 -0400
Subject: [PATCH] working on tests. WIP

---
 tests/crew_test.py | 14 +++++++++++---
 tests/task_test.py | 22 ++++++++++++++++------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/tests/crew_test.py b/tests/crew_test.py
index 5ea4347db..ac251caa6 100644
--- a/tests/crew_test.py
+++ b/tests/crew_test.py
@@ -702,7 +702,6 @@ def test_set_agents_step_callback():
         execute.return_value = "ok"
         crew.kickoff()
         assert researcher_agent.step_callback is not None
-        assert False
 
 
 def test_dont_set_agents_step_callback_if_already_set():
@@ -1037,7 +1036,7 @@ def test_crew_does_not_interpolate_without_inputs():
 
 
 def test_task_callback_on_crew():
-    from unittest.mock import patch
+    from unittest.mock import MagicMock, patch
 
     researcher_agent = Agent(
         role="Researcher",
@@ -1053,17 +1052,23 @@ def test_task_callback_on_crew():
         async_execution=True,
     )
 
+    mock_callback = MagicMock()
+
     crew = Crew(
         agents=[researcher_agent],
         process=Process.sequential,
         tasks=[list_ideas],
-        task_callback=lambda: None,
+        task_callback=mock_callback,
     )
 
     with patch.object(Agent, "execute_task") as execute:
         execute.return_value = "ok"
         crew.kickoff()
+
         assert list_ideas.callback is not None
+        mock_callback.assert_called_once()
+        args, _ = mock_callback.call_args
+        assert isinstance(args[0], TaskOutput)
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
@@ -1333,3 +1338,6 @@ def test_crew_train_error():
         assert "train() missing 1 required positional argument: 'n_iterations'" in str(
             e
         )
+
+
+# TODO: Test delegation in heirarchical process
diff --git a/tests/task_test.py b/tests/task_test.py
index 75190327e..16b75181b 100644
--- a/tests/task_test.py
+++ b/tests/task_test.py
@@ -1,7 +1,6 @@
 """Test Agent creation and execution basic functionality."""
 
 import json
-
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -81,7 +80,7 @@ def test_task_prompt_includes_expected_output():
 
     with patch.object(Agent, "execute_task") as execute:
         execute.return_value = "ok"
-        task.execute()
+        task.execute_sync()
         execute.assert_called_once_with(task=task, context=None, tools=[])
 
 
@@ -104,7 +103,7 @@ def test_task_callback():
 
     with patch.object(Agent, "execute_task") as execute:
         execute.return_value = "ok"
-        task.execute()
+        task.execute_sync()
         task_completed.assert_called_once_with(task.output)
 
 
@@ -127,7 +126,7 @@ def test_task_callback_returns_task_ouput():
 
     with patch.object(Agent, "execute_task") as execute:
         execute.return_value = "exported_ok"
-        task.execute()
+        task.execute_sync()
         # Ensure the callback is called with a TaskOutput object serialized to JSON
         task_completed.assert_called_once()
         callback_data = task_completed.call_args[0][0]
@@ -162,7 +161,7 @@ def test_execute_with_agent():
     )
 
     with patch.object(Agent, "execute_task", return_value="ok") as execute:
-        task.execute(agent=researcher)
+        task.execute_sync(agent=researcher)
         execute.assert_called_once_with(task=task, context=None, tools=[])
 
 
@@ -182,7 +181,7 @@ def test_async_execution():
     )
 
     with patch.object(Agent, "execute_task", return_value="ok") as execute:
-        task.execute(agent=researcher)
+        task.execute_async(agent=researcher)
         execute.assert_called_once_with(task=task, context=None, tools=[])
 
 
@@ -526,3 +525,14 @@ def test_interpolate_inputs():
         == "Give me a list of 5 interesting ideas about ML to explore for an article, what makes them unique and interesting."
     )
     assert task.expected_output == "Bullet point list of 5 interesting ideas about ML."
+
+
+"""
+TODO: TEST SYNC
+- Verify return type
+"""
+
+"""
+TODO: TEST ASYNC
+- Verify return type
+"""