1212# See the License for the specific language governing permissions and
1313# limitations under the License.
1414
15- import pytest
1615from unittest .mock import MagicMock
17- from google .adk .code_executors .unsafe_local_code_executor import UnsafeLocalCodeExecutor
18- from google .adk .code_executors .code_execution_utils import CodeExecutionInput , CodeExecutionResult
19- from google .adk .agents .invocation_context import InvocationContext
16+
2017from google .adk .agents .base_agent import BaseAgent
21- from google .adk .sessions .session import Session
18+ from google .adk .agents .invocation_context import InvocationContext
19+ from google .adk .code_executors .code_execution_utils import CodeExecutionInput
20+ from google .adk .code_executors .code_execution_utils import CodeExecutionResult
21+ from google .adk .code_executors .unsafe_local_code_executor import UnsafeLocalCodeExecutor
2222from google .adk .sessions .base_session_service import BaseSessionService
23+ from google .adk .sessions .session import Session
24+ import pytest
2325
2426
2527@pytest .fixture
2628def mock_invocation_context () -> InvocationContext :
27- """Provides a mock InvocationContext."""
28- mock_agent = MagicMock (spec = BaseAgent )
29- mock_session = MagicMock (spec = Session )
30- mock_session_service = MagicMock (spec = BaseSessionService )
31- return InvocationContext (
32- invocation_id = "test_invocation" ,
33- agent = mock_agent ,
34- session = mock_session ,
35- session_service = mock_session_service
36- )
29+ """Provides a mock InvocationContext."""
30+ mock_agent = MagicMock (spec = BaseAgent )
31+ mock_session = MagicMock (spec = Session )
32+ mock_session_service = MagicMock (spec = BaseSessionService )
33+ return InvocationContext (
34+ invocation_id = "test_invocation" ,
35+ agent = mock_agent ,
36+ session = mock_session ,
37+ session_service = mock_session_service ,
38+ )
3739
3840
3941class TestUnsafeLocalCodeExecutor :
4042
41- def test_init_default (self ):
42- executor = UnsafeLocalCodeExecutor ()
43- assert not executor .stateful
44- assert not executor .optimize_data_file
45-
46- def test_init_stateful_raises_error (self ):
47- with pytest .raises (ValueError , match = "Cannot set `stateful=True` in UnsafeLocalCodeExecutor." ):
48- UnsafeLocalCodeExecutor (stateful = True )
49-
50- def test_init_optimize_data_file_raises_error (self ):
51- with pytest .raises (ValueError , match = "Cannot set `optimize_data_file=True` in UnsafeLocalCodeExecutor." ):
52- UnsafeLocalCodeExecutor (optimize_data_file = True )
53-
54- def test_execute_code_simple_print (self , mock_invocation_context : InvocationContext ):
55- executor = UnsafeLocalCodeExecutor ()
56- code_input = CodeExecutionInput (code = 'print("hello world")' )
57- result = executor .execute_code (mock_invocation_context , code_input )
58-
59- assert isinstance (result , CodeExecutionResult )
60- assert result .stdout == "hello world\n "
61- assert result .stderr == ""
62- assert result .output_files == []
63-
64- def test_execute_code_with_error (self , mock_invocation_context : InvocationContext ):
65- executor = UnsafeLocalCodeExecutor ()
66- code_input = CodeExecutionInput (code = 'raise ValueError("Test error")' )
67- result = executor .execute_code (mock_invocation_context , code_input )
68-
69- assert isinstance (result , CodeExecutionResult )
70- assert result .stdout == ""
71- assert "Test error" in result .stderr
72- assert result .output_files == []
73-
74- def test_execute_code_variable_assignment (self , mock_invocation_context : InvocationContext ):
75- executor = UnsafeLocalCodeExecutor ()
76- code_input = CodeExecutionInput (code = 'x = 10\n print(x * 2)' )
77- result = executor .execute_code (mock_invocation_context , code_input )
78-
79- assert result .stdout == "20\n "
80- assert result .stderr == ""
81-
82- def test_execute_code_empty (self , mock_invocation_context : InvocationContext ):
83- executor = UnsafeLocalCodeExecutor ()
84- code_input = CodeExecutionInput (code = '' )
85- result = executor .execute_code (mock_invocation_context , code_input )
86- assert result .stdout == ""
87- assert result .stderr == ""
43+ def test_init_default (self ):
44+ executor = UnsafeLocalCodeExecutor ()
45+ assert not executor .stateful
46+ assert not executor .optimize_data_file
47+
48+ def test_init_stateful_raises_error (self ):
49+ with pytest .raises (
50+ ValueError ,
51+ match = "Cannot set `stateful=True` in UnsafeLocalCodeExecutor." ,
52+ ):
53+ UnsafeLocalCodeExecutor (stateful = True )
54+
55+ def test_init_optimize_data_file_raises_error (self ):
56+ with pytest .raises (
57+ ValueError ,
58+ match = (
59+ "Cannot set `optimize_data_file=True` in UnsafeLocalCodeExecutor."
60+ ),
61+ ):
62+ UnsafeLocalCodeExecutor (optimize_data_file = True )
63+
64+ def test_execute_code_simple_print (
65+ self , mock_invocation_context : InvocationContext
66+ ):
67+ executor = UnsafeLocalCodeExecutor ()
68+ code_input = CodeExecutionInput (code = 'print("hello world")' )
69+ result = executor .execute_code (mock_invocation_context , code_input )
70+
71+ assert isinstance (result , CodeExecutionResult )
72+ assert result .stdout == "hello world\n "
73+ assert result .stderr == ""
74+ assert result .output_files == []
75+
76+ def test_execute_code_with_error (
77+ self , mock_invocation_context : InvocationContext
78+ ):
79+ executor = UnsafeLocalCodeExecutor ()
80+ code_input = CodeExecutionInput (code = 'raise ValueError("Test error")' )
81+ result = executor .execute_code (mock_invocation_context , code_input )
82+
83+ assert isinstance (result , CodeExecutionResult )
84+ assert result .stdout == ""
85+ assert "Test error" in result .stderr
86+ assert result .output_files == []
87+
88+ def test_execute_code_variable_assignment (
89+ self , mock_invocation_context : InvocationContext
90+ ):
91+ executor = UnsafeLocalCodeExecutor ()
92+ code_input = CodeExecutionInput (code = "x = 10\n print(x * 2)" )
93+ result = executor .execute_code (mock_invocation_context , code_input )
94+
95+ assert result .stdout == "20\n "
96+ assert result .stderr == ""
97+
98+ def test_execute_code_empty (self , mock_invocation_context : InvocationContext ):
99+ executor = UnsafeLocalCodeExecutor ()
100+ code_input = CodeExecutionInput (code = "" )
101+ result = executor .execute_code (mock_invocation_context , code_input )
102+ assert result .stdout == ""
103+ assert result .stderr == ""
0 commit comments