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
|
import tempfile
import unittest
from unittest.mock import NonCallableMock, patch
import tools.jit.gen_unboxing as gen_unboxing
@patch("tools.jit.gen_unboxing.get_custom_build_selector")
@patch("tools.jit.gen_unboxing.parse_native_yaml")
@patch("tools.jit.gen_unboxing.make_file_manager")
@patch("tools.jit.gen_unboxing.gen_unboxing")
class TestGenUnboxing(unittest.TestCase):
def test_get_custom_build_selector_with_allowlist(
self,
mock_gen_unboxing: NonCallableMock,
mock_make_file_manager: NonCallableMock,
mock_parse_native_yaml: NonCallableMock,
mock_get_custom_build_selector: NonCallableMock,
) -> None:
args = ["--op-registration-allowlist=op1", "--op-selection-yaml-path=path2"]
gen_unboxing.main(args)
mock_get_custom_build_selector.assert_called_once_with(["op1"], "path2")
def test_get_custom_build_selector_with_allowlist_yaml(
self,
mock_gen_unboxing: NonCallableMock,
mock_make_file_manager: NonCallableMock,
mock_parse_native_yaml: NonCallableMock,
mock_get_custom_build_selector: NonCallableMock,
) -> None:
temp_file = tempfile.NamedTemporaryFile()
temp_file.write(b"- aten::add.Tensor")
temp_file.seek(0)
args = [
f"--TEST-ONLY-op-registration-allowlist-yaml-path={temp_file.name}",
"--op-selection-yaml-path=path2",
]
gen_unboxing.main(args)
mock_get_custom_build_selector.assert_called_once_with(
["aten::add.Tensor"], "path2"
)
temp_file.close()
def test_get_custom_build_selector_with_both_allowlist_and_yaml(
self,
mock_gen_unboxing: NonCallableMock,
mock_make_file_manager: NonCallableMock,
mock_parse_native_yaml: NonCallableMock,
mock_get_custom_build_selector: NonCallableMock,
) -> None:
temp_file = tempfile.NamedTemporaryFile()
temp_file.write(b"- aten::add.Tensor")
temp_file.seek(0)
args = [
"--op-registration-allowlist=op1",
"--TEST-ONLY-op-registration-allowlist-yaml-path={temp_file.name}",
"--op-selection-yaml-path=path2",
]
gen_unboxing.main(args)
mock_get_custom_build_selector.assert_called_once_with(["op1"], "path2")
temp_file.close()
if __name__ == "__main__":
unittest.main()
|
