File size: 1,286 Bytes
ed4d993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import unittest

from langchain_experimental.utilities import PythonREPL


class TestSanitizeInput(unittest.TestCase):
    def test_whitespace_removal(self) -> None:
        query = "   print('Hello, world!')   "
        sanitized_query = PythonREPL.sanitize_input(query)
        self.assertEqual(sanitized_query, "print('Hello, world!')")

    def test_python_removal(self) -> None:
        query = "python   print('Hello, world!')   "
        sanitized_query = PythonREPL.sanitize_input(query)
        self.assertEqual(sanitized_query, "print('Hello, world!')")

    def test_backtick_removal(self) -> None:
        query = "`print('Hello, world!')`"
        sanitized_query = PythonREPL.sanitize_input(query)
        self.assertEqual(sanitized_query, "print('Hello, world!')")

    def test_combined_removal(self) -> None:
        query = "  `python  print('Hello, world!')`  "
        sanitized_query = PythonREPL.sanitize_input(query)
        self.assertEqual(sanitized_query, "print('Hello, world!')")

    def test_mixed_case_removal(self) -> None:
        query = "  pYtHoN   print('Hello, world!')  "
        sanitized_query = PythonREPL.sanitize_input(query)
        self.assertEqual(sanitized_query, "print('Hello, world!')")


if __name__ == "__main__":
    unittest.main()