pytorch

1

2

3

4

5

6
from caffe2.python import context, test_util
7
from threading import Thread
8

9

10
class MyContext(context.Managed):
11
    pass
12

13
class DefaultMyContext(context.DefaultManaged):
14
    pass
15

16
class ChildMyContext(MyContext):
17
    pass
18

19

20
class TestContext(test_util.TestCase):
21
    def use_my_context(self):
22
        try:
23
            for _ in range(100):
24
                with MyContext() as a:
25
                    for _ in range(100):
26
                        self.assertTrue(MyContext.current() == a)
27
        except Exception as e:
28
            self._exceptions.append(e)
29

30
    def testMultiThreaded(self):
31
        threads = []
32
        self._exceptions = []
33
        for _ in range(8):
34
            thread = Thread(target=self.use_my_context)
35
            thread.start()
36
            threads.append(thread)
37
        for t in threads:
38
            t.join()
39
        for e in self._exceptions:
40
            raise e
41

42
    @MyContext()
43
    def testDecorator(self):
44
        self.assertIsNotNone(MyContext.current())
45

46
    def testNonDefaultCurrent(self):
47
        with self.assertRaises(AssertionError):
48
            MyContext.current()
49

50
        ctx = MyContext()
51
        self.assertEqual(MyContext.current(value=ctx), ctx)
52

53
        self.assertIsNone(MyContext.current(required=False))
54

55
    def testDefaultCurrent(self):
56
        self.assertIsInstance(DefaultMyContext.current(), DefaultMyContext)
57

58
    def testNestedContexts(self):
59
        with MyContext() as ctx1:
60
            with DefaultMyContext() as ctx2:
61
                self.assertEqual(DefaultMyContext.current(), ctx2)
62
                self.assertEqual(MyContext.current(), ctx1)
63

64
    def testChildClasses(self):
65
        with ChildMyContext() as ctx:
66
            self.assertEqual(ChildMyContext.current(), ctx)
67
            self.assertEqual(MyContext.current(), ctx)
68