1 from __future__ import absolute_import

     2 

     3 import silenttestrunner

     4 import unittest

     5 

     6 from mercurial.util import ctxmanager

     7 

     8 class contextmanager(object):

     9     def __init__(self, name, trace):

    10         self.name = name

    11         self.entered = False

    12         self.exited = False

    13         self.trace = trace

    14 

    15     def __enter__(self):

    16         self.entered = True

    17         self.trace(('enter', self.name))

    18         return self

    19 

    20     def __exit__(self, exc_type, exc_val, exc_tb):

    21         self.exited = exc_type, exc_val, exc_tb

    22         self.trace(('exit', self.name))

    23 

    24     def __repr__(self):

    25         return '<ctx %r>' % self.name

    26 

    27 class ctxerror(Exception):

    28     pass

    29 

    30 class raise_on_enter(contextmanager):

    31     def __enter__(self):

    32         self.trace(('raise', self.name))

    33         raise ctxerror(self.name)

    34 

    35 class raise_on_exit(contextmanager):

    36     def __exit__(self, exc_type, exc_val, exc_tb):

    37         self.trace(('raise', self.name))

    38         raise ctxerror(self.name)

    39 

    40 def ctxmgr(name, trace):

    41     return lambda: contextmanager(name, trace)

    42 

    43 class test_ctxmanager(unittest.TestCase):

    44     def test_basics(self):

    45         trace = []

    46         addtrace = trace.append

    47         with ctxmanager(ctxmgr('a', addtrace), ctxmgr('b', addtrace)) as c:

    48             a, b = c()

    49             c.atexit(addtrace, ('atexit', 'x'))

    50             c.atexit(addtrace, ('atexit', 'y'))

    51         self.assertEqual(trace, [('enter', 'a'), ('enter', 'b'),

    52                                  ('atexit', 'y'), ('atexit', 'x'),

    53                                  ('exit', 'b'), ('exit', 'a')])

    54 

    55     def test_raise_on_enter(self):

    56         trace = []

    57         addtrace = trace.append

    58         with self.assertRaises(ctxerror):

    59             with ctxmanager(ctxmgr('a', addtrace),

    60                            lambda: raise_on_enter('b', addtrace)) as c:

    61                 c()

    62                 addtrace('unreachable')

    63         self.assertEqual(trace, [('enter', 'a'), ('raise', 'b'), ('exit', 'a')])

    64 

    65     def test_raise_on_exit(self):

    66         trace = []

    67         addtrace = trace.append

    68         with self.assertRaises(ctxerror):

    69             with ctxmanager(ctxmgr('a', addtrace),

    70                            lambda: raise_on_exit('b', addtrace)) as c:

    71                 c()

    72                 addtrace('running')

    73         self.assertEqual(trace, [('enter', 'a'), ('enter', 'b'), 'running',

    74                                  ('raise', 'b'), ('exit', 'a')])

    75 

    76 if __name__ == '__main__':

    77     silenttestrunner.main(__name__)