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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
# Owner(s): ["oncall: package/deploy"]
from io import BytesIO
from torch.package import PackageExporter
from torch.testing._internal.common_utils import run_tests
try:
from .common import PackageTestCase
except ImportError:
# Support the case where we run this file directly.
from common import PackageTestCase
class TestDependencyHooks(PackageTestCase):
"""Dependency management hooks API tests.
- register_mock_hook()
- register_extern_hook()
"""
def test_single_hook(self):
buffer = BytesIO()
my_externs = set()
def my_extern_hook(package_exporter, module_name):
my_externs.add(module_name)
with PackageExporter(buffer) as exporter:
exporter.extern(["package_a.subpackage", "module_a"])
exporter.register_extern_hook(my_extern_hook)
exporter.save_source_string("foo", "import module_a")
self.assertEqual(my_externs, set(["module_a"]))
def test_multiple_extern_hooks(self):
buffer = BytesIO()
my_externs = set()
def my_extern_hook(package_exporter, module_name):
my_externs.add(module_name)
# This also checks ordering, since `remove()` will fail if the value is not in the set.
def my_extern_hook2(package_exporter, module_name):
my_externs.remove(module_name)
with PackageExporter(buffer) as exporter:
exporter.extern(["package_a.subpackage", "module_a"])
exporter.register_extern_hook(my_extern_hook)
exporter.register_extern_hook(my_extern_hook2)
exporter.save_source_string("foo", "import module_a")
self.assertEqual(my_externs, set())
def test_multiple_mock_hooks(self):
buffer = BytesIO()
my_mocks = set()
def my_mock_hook(package_exporter, module_name):
my_mocks.add(module_name)
# This also checks ordering, since `remove()` will fail if the value is not in the set.
def my_mock_hook2(package_exporter, module_name):
my_mocks.remove(module_name)
with PackageExporter(buffer) as exporter:
exporter.mock(["package_a.subpackage", "module_a"])
exporter.register_mock_hook(my_mock_hook)
exporter.register_mock_hook(my_mock_hook2)
exporter.save_source_string("foo", "import module_a")
self.assertEqual(my_mocks, set())
def test_remove_hooks(self):
buffer = BytesIO()
my_externs = set()
my_externs2 = set()
def my_extern_hook(package_exporter, module_name):
my_externs.add(module_name)
def my_extern_hook2(package_exporter, module_name):
my_externs2.add(module_name)
with PackageExporter(buffer) as exporter:
exporter.extern(["package_a.subpackage", "module_a"])
handle = exporter.register_extern_hook(my_extern_hook)
exporter.register_extern_hook(my_extern_hook2)
handle.remove()
exporter.save_source_string("foo", "import module_a")
self.assertEqual(my_externs, set())
self.assertEqual(my_externs2, set(["module_a"]))
def test_extern_and_mock_hook(self):
buffer = BytesIO()
my_externs = set()
my_mocks = set()
def my_extern_hook(package_exporter, module_name):
my_externs.add(module_name)
def my_mock_hook(package_exporter, module_name):
my_mocks.add(module_name)
with PackageExporter(buffer) as exporter:
exporter.extern("module_a")
exporter.mock("package_a")
exporter.register_extern_hook(my_extern_hook)
exporter.register_mock_hook(my_mock_hook)
exporter.save_source_string("foo", "import module_a; import package_a")
self.assertEqual(my_externs, set(["module_a"]))
self.assertEqual(my_mocks, set(["package_a"]))
if __name__ == "__main__":
run_tests()
|
