add back in guardrails

tests/test_task_guardrails.py

@@ -1,129 +1,129 @@
-# """Tests for task guardrails functionality."""
+"""Tests for task guardrails functionality."""
 
-# from unittest.mock import Mock
+from unittest.mock import Mock
 
-# import pytest
+import pytest
 
-# from crewai.task import Task
+from crewai.task import Task
-# from crewai.tasks.task_output import TaskOutput
+from crewai.tasks.task_output import TaskOutput
 
 
-# def test_task_without_guardrail():
+def test_task_without_guardrail():
-#     """Test that tasks work normally without guardrails (backward compatibility)."""
+    """Test that tasks work normally without guardrails (backward compatibility)."""
-#     agent = Mock()
+    agent = Mock()
-#     agent.role = "test_agent"
+    agent.role = "test_agent"
-#     agent.execute_task.return_value = "test result"
+    agent.execute_task.return_value = "test result"
-#     agent.crew = None
+    agent.crew = None
 
-#     task = Task(description="Test task", expected_output="Output")
+    task = Task(description="Test task", expected_output="Output")
 
-#     result = task.execute_sync(agent=agent)
+    result = task.execute_sync(agent=agent)
-#     assert isinstance(result, TaskOutput)
+    assert isinstance(result, TaskOutput)
-#     assert result.raw == "test result"
+    assert result.raw == "test result"
 
 
-# def test_task_with_successful_guardrail():
+def test_task_with_successful_guardrail():
-#     """Test that successful guardrail validation passes transformed result."""
+    """Test that successful guardrail validation passes transformed result."""
 
-#     def guardrail(result: TaskOutput):
+    def guardrail(result: TaskOutput):
-#         return (True, result.raw.upper())
+        return (True, result.raw.upper())
 
-#     agent = Mock()
+    agent = Mock()
-#     agent.role = "test_agent"
+    agent.role = "test_agent"
-#     agent.execute_task.return_value = "test result"
+    agent.execute_task.return_value = "test result"
-#     agent.crew = None
+    agent.crew = None
 
-#     task = Task(description="Test task", expected_output="Output", guardrail=guardrail)
+    task = Task(description="Test task", expected_output="Output", guardrail=guardrail)
 
-#     result = task.execute_sync(agent=agent)
+    result = task.execute_sync(agent=agent)
-#     assert isinstance(result, TaskOutput)
+    assert isinstance(result, TaskOutput)
-#     assert result.raw == "TEST RESULT"
+    assert result.raw == "TEST RESULT"
 
 
-# def test_task_with_failing_guardrail():
+def test_task_with_failing_guardrail():
-#     """Test that failing guardrail triggers retry with error context."""
+    """Test that failing guardrail triggers retry with error context."""
 
-#     def guardrail(result: TaskOutput):
+    def guardrail(result: TaskOutput):
-#         return (False, "Invalid format")
+        return (False, "Invalid format")
 
-#     agent = Mock()
+    agent = Mock()
-#     agent.role = "test_agent"
+    agent.role = "test_agent"
-#     agent.execute_task.side_effect = ["bad result", "good result"]
+    agent.execute_task.side_effect = ["bad result", "good result"]
-#     agent.crew = None
+    agent.crew = None
 
-#     task = Task(
+    task = Task(
-#         description="Test task",
+        description="Test task",
-#         expected_output="Output",
+        expected_output="Output",
-#         guardrail=guardrail,
+        guardrail=guardrail,
-#         max_retries=1,
+        max_retries=1,
-#     )
+    )
 
-#     # First execution fails guardrail, second succeeds
+    # First execution fails guardrail, second succeeds
-#     agent.execute_task.side_effect = ["bad result", "good result"]
+    agent.execute_task.side_effect = ["bad result", "good result"]
-#     with pytest.raises(Exception) as exc_info:
+    with pytest.raises(Exception) as exc_info:
-#         task.execute_sync(agent=agent)
+        task.execute_sync(agent=agent)
 
-#     assert "Task failed guardrail validation" in str(exc_info.value)
+    assert "Task failed guardrail validation" in str(exc_info.value)
-#     assert task.retry_count == 1
+    assert task.retry_count == 1
 
 
-# def test_task_with_guardrail_retries():
+def test_task_with_guardrail_retries():
-#     """Test that guardrail respects max_retries configuration."""
+    """Test that guardrail respects max_retries configuration."""
 
-#     def guardrail(result: TaskOutput):
+    def guardrail(result: TaskOutput):
-#         return (False, "Invalid format")
+        return (False, "Invalid format")
 
-#     agent = Mock()
+    agent = Mock()
-#     agent.role = "test_agent"
+    agent.role = "test_agent"
-#     agent.execute_task.return_value = "bad result"
+    agent.execute_task.return_value = "bad result"
-#     agent.crew = None
+    agent.crew = None
 
-#     task = Task(
+    task = Task(
-#         description="Test task",
+        description="Test task",
-#         expected_output="Output",
+        expected_output="Output",
-#         guardrail=guardrail,
+        guardrail=guardrail,
-#         max_retries=2,
+        max_retries=2,
-#     )
+    )
 
-#     with pytest.raises(Exception) as exc_info:
+    with pytest.raises(Exception) as exc_info:
-#         task.execute_sync(agent=agent)
+        task.execute_sync(agent=agent)
 
-#     assert task.retry_count == 2
+    assert task.retry_count == 2
-#     assert "Task failed guardrail validation after 2 retries" in str(exc_info.value)
+    assert "Task failed guardrail validation after 2 retries" in str(exc_info.value)
-#     assert "Invalid format" in str(exc_info.value)
+    assert "Invalid format" in str(exc_info.value)
 
 
-# def test_guardrail_error_in_context():
+def test_guardrail_error_in_context():
-#     """Test that guardrail error is passed in context for retry."""
+    """Test that guardrail error is passed in context for retry."""
 
-#     def guardrail(result: TaskOutput):
+    def guardrail(result: TaskOutput):
-#         return (False, "Expected JSON, got string")
+        return (False, "Expected JSON, got string")
 
-#     agent = Mock()
+    agent = Mock()
-#     agent.role = "test_agent"
+    agent.role = "test_agent"
-#     agent.crew = None
+    agent.crew = None
 
-#     task = Task(
+    task = Task(
-#         description="Test task",
+        description="Test task",
-#         expected_output="Output",
+        expected_output="Output",
-#         guardrail=guardrail,
+        guardrail=guardrail,
-#         max_retries=1,
+        max_retries=1,
-#     )
+    )
 
-#     # Mock execute_task to succeed on second attempt
+    # Mock execute_task to succeed on second attempt
-#     first_call = True
+    first_call = True
 
-#     def execute_task(task, context, tools):
+    def execute_task(task, context, tools):
-#         nonlocal first_call
+        nonlocal first_call
-#         if first_call:
+        if first_call:
-#             first_call = False
+            first_call = False
-#             return "invalid"
+            return "invalid"
-#         return '{"valid": "json"}'
+        return '{"valid": "json"}'
 
-#     agent.execute_task.side_effect = execute_task
+    agent.execute_task.side_effect = execute_task
 
-#     with pytest.raises(Exception) as exc_info:
+    with pytest.raises(Exception) as exc_info:
-#         task.execute_sync(agent=agent)
+        task.execute_sync(agent=agent)
 
-#     assert "Task failed guardrail validation" in str(exc_info.value)
+    assert "Task failed guardrail validation" in str(exc_info.value)
-#     assert "Expected JSON, got string" in str(exc_info.value)
+    assert "Expected JSON, got string" in str(exc_info.value)