Gracefully terminate a Python program with an infinite loop (and test it)
I've recently been writing long-running microservices in Python, all of which contain an infinite run loop:
def run():
while True:
# Do a thing
# Do another thing
# Do a third thing
print("Did things")
At any time, the program could terminate midway through its operations due to an external interruption. For some applications, this may be okay. Other times, we may want to complete a loop iteration before stopping the program.
To gracefully terminate the infinite loop, we can create a flag and use Python's signal module.
# graceful_termination.py
import signal
import time
run_flag = True
def signal_handler(sig, frame):
sig_enum = signal.Signals(sig)
print("Program received the signal: ", sig_enum.name)
global run_flag
run_flag = False
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
def run():
while run_flag:
print("Press Ctrl+C")
print("Sleeping...")
time.sleep(0.5)
print("Done sleeping")
if __name__ == "__main__":
run()
print("Exiting...")
The handler captures the SIGINT and SIGTERM interruptions1 and disables the global run_flag.
If the program is in the middle of a loop iteration, all the steps are completed, then the flag is checked, the condition fails, and the program quits gracefully.
Here's what the output looks like:
$ python graceful_termination.py
Press Ctrl+C
Sleeping...
Done sleeping
Press Ctrl+C
Sleeping...
Done sleeping
Press Ctrl+C
Sleeping...
^CProgram received the signal: SIGINT
Done sleeping
Exiting...
Neat!
Now, what if I'm stubborn about test coverage and want to test everything, including the infinitely loopy run() function?
We can do that by mocking the run_flag.
# test_graceful_termination.py
from unittest.mock import patch
from graceful_termination import run
def test_infinite_loop_run_func():
with patch("graceful_termination.run_flag") as run_flag_mock:
run_flag_mock.__bool__.side_effect = [True, True, False] # Run the loop twice
with patch("time.sleep", return_value=None) as sleep_mock:
run()
assert sleep_mock.call_count == 2 # Exactly 2 iterations of the loop occurred.
The side_effect specifies the value to be returned by the mock object each time it is accessed.
It returns True twice, then False.
By mocking time.sleep, we can verify that the sleep function was called exactly twice.