mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-12 01:28:30 +00:00
229bbd9bbe
- Add execution_image field to Agent class for specifying custom Docker images - Modify get_code_execution_tools() to pass custom image to CodeInterpreterTool - Add comprehensive tests for the new functionality - Update documentation with usage examples - Fixes #2933: Allow customizing CrewAI code executor Docker image Co-Authored-By: João <joao@crewai.com>
83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
from crewai import Agent
|
|
|
|
|
|
def test_agent_with_custom_execution_image():
|
|
"""Test that Agent can be created with custom execution image."""
|
|
agent = Agent(
|
|
role="Test Agent",
|
|
goal="Test goal",
|
|
backstory="Test backstory",
|
|
allow_code_execution=True,
|
|
execution_image="my-custom-image:latest"
|
|
)
|
|
|
|
assert agent.execution_image == "my-custom-image:latest"
|
|
|
|
|
|
def test_agent_without_custom_execution_image():
|
|
"""Test that Agent works without custom execution image (default behavior)."""
|
|
agent = Agent(
|
|
role="Test Agent",
|
|
goal="Test goal",
|
|
backstory="Test backstory",
|
|
allow_code_execution=True
|
|
)
|
|
|
|
assert agent.execution_image is None
|
|
|
|
|
|
@patch('crewai.agent.CodeInterpreterTool')
|
|
def test_get_code_execution_tools_with_custom_image(mock_code_interpreter):
|
|
"""Test that get_code_execution_tools passes custom image to CodeInterpreterTool."""
|
|
agent = Agent(
|
|
role="Test Agent",
|
|
goal="Test goal",
|
|
backstory="Test backstory",
|
|
allow_code_execution=True,
|
|
execution_image="my-custom-image:latest"
|
|
)
|
|
|
|
tools = agent.get_code_execution_tools()
|
|
|
|
mock_code_interpreter.assert_called_once_with(
|
|
unsafe_mode=False,
|
|
default_image_tag="my-custom-image:latest"
|
|
)
|
|
|
|
|
|
@patch('crewai.agent.CodeInterpreterTool')
|
|
def test_get_code_execution_tools_without_custom_image(mock_code_interpreter):
|
|
"""Test that get_code_execution_tools works without custom image."""
|
|
agent = Agent(
|
|
role="Test Agent",
|
|
goal="Test goal",
|
|
backstory="Test backstory",
|
|
allow_code_execution=True
|
|
)
|
|
|
|
tools = agent.get_code_execution_tools()
|
|
|
|
mock_code_interpreter.assert_called_once_with(unsafe_mode=False)
|
|
|
|
|
|
@patch('crewai.agent.CodeInterpreterTool')
|
|
def test_get_code_execution_tools_with_unsafe_mode_and_custom_image(mock_code_interpreter):
|
|
"""Test that both unsafe_mode and custom image work together."""
|
|
agent = Agent(
|
|
role="Test Agent",
|
|
goal="Test goal",
|
|
backstory="Test backstory",
|
|
allow_code_execution=True,
|
|
code_execution_mode="unsafe",
|
|
execution_image="my-custom-image:latest"
|
|
)
|
|
|
|
tools = agent.get_code_execution_tools()
|
|
|
|
mock_code_interpreter.assert_called_once_with(
|
|
unsafe_mode=True,
|
|
default_image_tag="my-custom-image:latest"
|
|
)
|