Skip to content

Instantly share code, notes, and snippets.

@itsthejoker
Created August 23, 2016 18:41
Show Gist options
  • Save itsthejoker/6b497f2098916cefd8a8e2f9a1ff7b5d to your computer and use it in GitHub Desktop.
Save itsthejoker/6b497f2098916cefd8a8e2f9a1ff7b5d to your computer and use it in GitHub Desktop.
import errno
import os
import signal
import unittest
import time
class graceful_interrupt_handler(object):
'''
Usage:
with graceful_interrupt_handler as handler:
do_stuff()
if handler.interrupted:
do_more_stuff()
handler.interrupted is called immediately upon receiving the termination
signal, so do_more_stuff() in the above example should be something that
allows do_stuff to finish cleanly. The script will exit after
handler.interrupted finishes unless you wrap calls in itself, at which
point it will require as many consecutive calls to kill as you wrap.
A fully tested suite based off http://stackoverflow.com/a/10972804/2638784
'''
def __init__(self, signals=(signal.SIGINT, signal.SIGTERM)):
self.signals = signals
self.original_handlers = {}
def __enter__(self):
self.interrupted = False
self.released = False
for sig in self.signals:
self.original_handlers[sig] = signal.getsignal(sig)
signal.signal(sig, self.handler)
return self
def handler(self, signum, frame):
self.release()
self.interrupted = True
def __exit__(self, type, value, tb):
self.release()
def release(self):
if self.released:
return False
for sig in self.signals:
signal.signal(sig, self.original_handlers[sig])
self.released = True
return True
#################################################
# TESTING
#################################################
TESTFN = "$-testfile"
TESTFN2 = "$-testfile2"
POSIX = os.name == 'posix'
WINDOWS = os.name == 'nt'
test_files = []
def safe_remove(file):
"Convenience function for removing temporary test files"
try:
os.remove(file)
except OSError as err:
if err.errno != errno.ENOENT:
raise
@unittest.skipIf(WINDOWS, "Unix-only test!")
class GracefulInterruptHandlerTestCase(unittest.TestCase):
def test_exit_cleanly(self):
# Make sure handler is called on SIGTERM
pid = os.fork()
if pid:
# we're the parent process
time.sleep(1)
os.kill(pid, signal.SIGTERM)
time.sleep(1)
self.assertTrue(os.path.exists(TESTFN))
else:
# we're the child process
with graceful_interrupt_handler() as h:
time.sleep(5)
if h.interrupted:
open(TESTFN, 'wb').write('yo')
def test_exit_cleanly_with_SIGKILL(self):
# Make sure handler is not called on SIGKILL
pid = os.fork()
if pid:
# we're the parent process
time.sleep(1)
os.kill(pid, signal.SIGKILL)
time.sleep(1)
self.assertFalse(os.path.exists(TESTFN2))
else:
# we're the child process
with graceful_interrupt_handler() as h:
time.sleep(5)
if h.interrupted:
# will never get created
open(TESTFN2, 'wb').write('yo')
if __name__ == '__main__':
safe_remove(TESTFN)
unittest.main(verbosity=2)
safe_remove(TESTFN)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment