-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathtest_function_cli.py
More file actions
92 lines (66 loc) · 2.86 KB
/
test_function_cli.py
File metadata and controls
92 lines (66 loc) · 2.86 KB
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import unittest
from polyapi.function_cli import _parse_code
SIMPLE_CODE = """
def foobar(n: int) -> int:
return 9
"""
COMPLEX_RETURN_TYPE = """
from typing_extensions import TypedDict
class Barbar(TypedDict):
count: int
def foobar(n: int) -> Barbar:
return Barbar(count=n)
"""
LIST_COMPLEX_RETURN_TYPE = """
from typing import List
from typing_extensions import TypedDict
class Barbar(TypedDict):
count: int
def foobar(n: int) -> List[Barbar]:
return [Barbar(count=n)]
"""
COMPLEX_ARG_TYPE = """
from typing_extensions import TypedDict
class Barbar(TypedDict):
count: int
def foobar(n: Barbar) -> int:
return 7
"""
class T(unittest.TestCase):
def test_simple_types(self):
args, return_type, return_type_schema, additional_requirements = _parse_code(SIMPLE_CODE, "foobar")
self.assertEqual(len(args), 1)
self.assertEqual(args[0], {"key": "n", "name": "n", "type": "integer"})
self.assertEqual(return_type, "integer")
self.assertIsNone(return_type_schema)
self.assertEqual(additional_requirements, [])
def test_complex_return_type(self):
args, return_type, return_type_schema, _ = _parse_code(COMPLEX_RETURN_TYPE, "foobar")
self.assertEqual(len(args), 1)
self.assertEqual(args[0], {"key": "n", "name": "n", "type": "integer"})
self.assertEqual(return_type, "Barbar")
self.assertEqual(return_type_schema['title'], "Barbar")
def test_complex_arg_type(self):
args, return_type, return_type_schema, _ = _parse_code(COMPLEX_ARG_TYPE, "foobar")
self.assertEqual(len(args), 1)
self.assertEqual(args[0]["type"], "Barbar")
self.assertEqual(return_type, "integer")
self.assertIsNone(return_type_schema)
def test_list_complex_return_type(self):
args, return_type, return_type_schema, _ = _parse_code(LIST_COMPLEX_RETURN_TYPE, "foobar")
self.assertEqual(len(args), 1)
self.assertEqual(args[0], {"key": "n", "name": "n", "type": "integer"})
self.assertEqual(return_type, "object")
self.assertEqual(return_type_schema["items"]['title'], "Barbar")
def test_parse_import_basic(self):
code = "import flask\n\n\ndef foobar(n: int) -> int:\n return 9\n"
_, _, _, additional_requirements = _parse_code(code, "foobar")
self.assertEqual(additional_requirements, ["flask"])
def test_parse_import_from(self):
code = "from flask import Request, Response\n\n\ndef foobar(n: int) -> int:\n return 9\n"
_, _, _, additional_requirements = _parse_code(code, "foobar")
self.assertEqual(additional_requirements, ["flask"])
def test_parse_import_base(self):
code = "import requests\n\n\ndef foobar(n: int) -> int:\n return 9\n"
_, _, _, additional_requirements = _parse_code(code, "foobar")
self.assertEqual(additional_requirements, [])