test_win32_shim.py revision d3907f0d
1from __future__ import print_function
2
3import os
4
5from functools import wraps
6from zmq.tests import BaseZMQTestCase
7from zmq.utils.win32 import allow_interrupt
8
9
10def count_calls(f):
11    @wraps(f)
12    def _(*args, **kwds):
13        try:
14            return f(*args, **kwds)
15        finally:
16            _.__calls__ += 1
17    _.__calls__ = 0
18    return _
19
20
21class TestWindowsConsoleControlHandler(BaseZMQTestCase):
22
23    def test_handler(self):
24        @count_calls
25        def interrupt_polling():
26            print('Caught CTRL-C!')
27
28        if os.name == 'nt':
29            from ctypes import windll
30            from ctypes.wintypes import BOOL, DWORD
31
32            kernel32 = windll.LoadLibrary('kernel32')
33
34            # <http://msdn.microsoft.com/en-us/library/ms683155.aspx>
35            GenerateConsoleCtrlEvent = kernel32.GenerateConsoleCtrlEvent
36            GenerateConsoleCtrlEvent.argtypes = (DWORD, DWORD)
37            GenerateConsoleCtrlEvent.restype = BOOL
38
39            try:
40                # Simulate CTRL-C event while handler is active.
41                with allow_interrupt(interrupt_polling):
42                    result = GenerateConsoleCtrlEvent(0, 0)
43                    if result == 0:
44                        raise WindowsError
45            except KeyboardInterrupt:
46                pass
47            else:
48                self.fail('Expecting `KeyboardInterrupt` exception!')
49
50            # Make sure our handler was called.
51            self.assertEqual(interrupt_polling.__calls__, 1)
52        else:
53            # On non-Windows systems, this utility is just a no-op!
54            with allow_interrupt(interrupt_polling):
55                pass
56            self.assertEqual(interrupt_polling.__calls__, 0)
57