Handle both SIGINT and SIGTERM for docker-compose up.

Signed-off-by: Daniel Nephin <[email protected]>

compose/cli/main.py

@@ -658,17 +658,24 @@ def build_log_printer(containers, service_names, monochrome):
 
 def attach_to_logs(project, log_printer, service_names, timeout):
     print("Attaching to", list_containers(log_printer.containers))
-    try:
-        log_printer.run()
-    finally:
-        def handler(signal, frame):
-            project.kill(service_names=service_names)
-            sys.exit(0)
-        signal.signal(signal.SIGINT, handler)
 
+    def force_shutdown(signal, frame):
+        project.kill(service_names=service_names)
+        sys.exit(2)
+
+    def shutdown(signal, frame):
+        set_signal_handler(force_shutdown)
         print("Gracefully stopping... (press Ctrl+C again to force)")
         project.stop(service_names=service_names, timeout=timeout)
 
+    set_signal_handler(shutdown)
+    log_printer.run()
+
+
+def set_signal_handler(handler):
+    signal.signal(signal.SIGINT, handler)
+    signal.signal(signal.SIGTERM, handler)
+
 
 def list_containers(containers):
     return ", ".join(c.name for c in containers)

tests/acceptance/cli_test.py

@@ -2,7 +2,9 @@
 
 import os
 import shlex
+import signal
 import subprocess
+import time
 from collections import namedtuple
 from operator import attrgetter
 
@@ -20,6 +22,45 @@
 BUILD_PULL_TEXT = 'Status: Image is up to date for busybox:latest'
 
 
+def start_process(base_dir, options):
+    proc = subprocess.Popen(
+        ['docker-compose'] + options,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        cwd=base_dir)
+    print("Running process: %s" % proc.pid)
+    return proc
+
+
+def wait_on_process(proc, returncode=0):
+    stdout, stderr = proc.communicate()
+    if proc.returncode != returncode:
+        print(stderr)
+        assert proc.returncode == returncode
+    return ProcessResult(stdout.decode('utf-8'), stderr.decode('utf-8'))
+
+
+def wait_on_condition(condition, delay=0.1, timeout=5):
+    start_time = time.time()
+    while not condition():
+        if time.time() - start_time > timeout:
+            raise AssertionError("Timeout: %s" % condition)
+        time.sleep(delay)
+
+
+class ContainerCountCondition(object):
+
+    def __init__(self, project, expected):
+        self.project = project
+        self.expected = expected
+
+    def __call__(self):
+        return len(self.project.containers()) == self.expected
+
+    def __str__(self):
+        return "waiting for counter count == %s" % self.expected
+
+
 class CLITestCase(DockerClientTestCase):
 
     def setUp(self):
@@ -42,17 +83,8 @@ def project(self):
 
     def dispatch(self, options, project_options=None, returncode=0):
         project_options = project_options or []
-        proc = subprocess.Popen(
-            ['docker-compose'] + project_options + options,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=self.base_dir)
-        print("Running process: %s" % proc.pid)
-        stdout, stderr = proc.communicate()
-        if proc.returncode != returncode:
-            print(stderr)
-            assert proc.returncode == returncode
-        return ProcessResult(stdout.decode('utf-8'), stderr.decode('utf-8'))
+        proc = start_process(self.base_dir, project_options + options)
+        return wait_on_process(proc, returncode=returncode)
 
     def test_help(self):
         old_base_dir = self.base_dir
@@ -291,7 +323,7 @@ def test_up_with_force_recreate_and_no_recreate(self):
             returncode=1)
 
     def test_up_with_timeout(self):
-        self.dispatch(['up', '-d', '-t', '1'], None)
+        self.dispatch(['up', '-d', '-t', '1'])
         service = self.project.get_service('simple')
         another = self.project.get_service('another')
         self.assertEqual(len(service.containers()), 1)
@@ -303,6 +335,20 @@ def test_up_with_timeout(self):
         self.assertFalse(config['AttachStdout'])
         self.assertFalse(config['AttachStdin'])
 
+    def test_up_handles_sigint(self):
+        proc = start_process(self.base_dir, ['up', '-t', '2'])
+        wait_on_condition(ContainerCountCondition(self.project, 2))
+
+        os.kill(proc.pid, signal.SIGINT)
+        wait_on_condition(ContainerCountCondition(self.project, 0))
+
+    def test_up_handles_sigterm(self):
+        proc = start_process(self.base_dir, ['up', '-t', '2'])
+        wait_on_condition(ContainerCountCondition(self.project, 2))
+
+        os.kill(proc.pid, signal.SIGTERM)
+        wait_on_condition(ContainerCountCondition(self.project, 0))
+
     def test_run_service_without_links(self):
         self.base_dir = 'tests/fixtures/links-composefile'
         self.dispatch(['run', 'console', '/bin/true'])

tests/unit/cli/main_test.py

@@ -57,11 +57,11 @@ def test_attach_to_logs(self):
         with mock.patch('compose.cli.main.signal', autospec=True) as mock_signal:
             attach_to_logs(project, log_printer, service_names, timeout)
 
-        mock_signal.signal.assert_called_once_with(mock_signal.SIGINT, mock.ANY)
+        assert mock_signal.signal.mock_calls == [
+            mock.call(mock_signal.SIGINT, mock.ANY),
+            mock.call(mock_signal.SIGTERM, mock.ANY),
+        ]
         log_printer.run.assert_called_once_with()
-        project.stop.assert_called_once_with(
-            service_names=service_names,
-            timeout=timeout)
 
 
 class SetupConsoleHandlerTestCase(unittest.TestCase):