File size: 832 Bytes
ac819bc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from unittest import TestCase

from lagent.actions.python_interpreter import PythonInterpreter
from lagent.schema import ActionStatusCode


class TestPythonInterpreter(TestCase):

    def test_python_executor(self):
        python_executor = PythonInterpreter()
        tool_return = python_executor(
            '```python\ndef solution():\n    return 1\n```')
        self.assertEqual(tool_return.state, ActionStatusCode.SUCCESS)
        self.assertDictEqual(tool_return.result, dict(text='1'))

    def test_timeout(self):
        python_executor = PythonInterpreter(timeout=2)
        tool_return = python_executor(
            '```python\ndef solution():\n    while True:\n        pass\n```')
        self.assertEqual(tool_return.state, ActionStatusCode.API_ERROR)
        self.assertIn('FunctionTimedOut', tool_return.errmsg)