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)