From 778c213dfc9c65ac27d4f527018eb66d841d890e Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@docker.com>
Date: Mon, 28 Dec 2015 16:57:55 -0500
Subject: [PATCH] Fix signal handlers by moving shutdown logic out of handler.

Signed-off-by: Daniel Nephin <dnephin@docker.com>
---
 compose/cli/main.py          | 45 +++++++++++++++---------------------
 compose/cli/signals.py       | 18 +++++++++++++++
 tests/acceptance/cli_test.py |  3 ++-
 tests/unit/cli/main_test.py  |  2 +-
 4 files changed, 40 insertions(+), 28 deletions(-)
 create mode 100644 compose/cli/signals.py

diff --git a/compose/cli/main.py b/compose/cli/main.py
index 4a766133f..e360dddd1 100644
--- a/compose/cli/main.py
+++ b/compose/cli/main.py
@@ -3,7 +3,6 @@ from __future__ import unicode_literals
 
 import logging
 import re
-import signal
 import sys
 from inspect import getdoc
 from operator import attrgetter
@@ -12,6 +11,7 @@ import yaml
 from docker.errors import APIError
 from requests.exceptions import ReadTimeout
 
+from . import signals
 from .. import __version__
 from ..config import config
 from ..config import ConfigurationError
@@ -655,20 +655,19 @@ def run_one_off_container(container_options, project, service, options):
         if options['--rm']:
             project.client.remove_container(container.id, force=True)
 
-    def force_shutdown(signal, frame):
+    signals.set_signal_handler_to_shutdown()
+    try:
+        try:
+            dockerpty.start(project.client, container.id, interactive=not options['-T'])
+            exit_code = container.wait()
+        except signals.ShutdownException:
+            project.client.stop(container.id)
+            exit_code = 1
+    except signals.ShutdownException:
         project.client.kill(container.id)
         remove_container(force=True)
         sys.exit(2)
 
-    def shutdown(signal, frame):
-        set_signal_handler(force_shutdown)
-        project.client.stop(container.id)
-        remove_container()
-        sys.exit(1)
-
-    set_signal_handler(shutdown)
-    dockerpty.start(project.client, container.id, interactive=not options['-T'])
-    exit_code = container.wait()
     remove_container()
     sys.exit(exit_code)
 
@@ -683,25 +682,19 @@ 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))
+    signals.set_signal_handler_to_shutdown()
 
-    def force_shutdown(signal, frame):
+    try:
+        try:
+            log_printer.run()
+        except signals.ShutdownException:
+            print("Gracefully stopping... (press Ctrl+C again to force)")
+            project.stop(service_names=service_names, timeout=timeout)
+    except signals.ShutdownException:
         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)
-
-    print("Attaching to", list_containers(log_printer.containers))
-    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)
diff --git a/compose/cli/signals.py b/compose/cli/signals.py
new file mode 100644
index 000000000..38474ba44
--- /dev/null
+++ b/compose/cli/signals.py
@@ -0,0 +1,18 @@
+import signal
+
+
+class ShutdownException(Exception):
+    pass
+
+
+def shutdown(signal, frame):
+    raise ShutdownException()
+
+
+def set_signal_handler(handler):
+    signal.signal(signal.SIGINT, handler)
+    signal.signal(signal.SIGTERM, handler)
+
+
+def set_signal_handler_to_shutdown():
+    set_signal_handler(shutdown)
diff --git a/tests/acceptance/cli_test.py b/tests/acceptance/cli_test.py
index e6fa38a8b..694f8583f 100644
--- a/tests/acceptance/cli_test.py
+++ b/tests/acceptance/cli_test.py
@@ -86,7 +86,8 @@ class ContainerStateCondition(object):
             return False
 
     def __str__(self):
-        return "waiting for container to have state %s" % self.expected
+        state = 'running' if self.running else 'stopped'
+        return "waiting for container to be %s" % state
 
 
 class CLITestCase(DockerClientTestCase):
diff --git a/tests/unit/cli/main_test.py b/tests/unit/cli/main_test.py
index db37ac1af..b63ac7468 100644
--- a/tests/unit/cli/main_test.py
+++ b/tests/unit/cli/main_test.py
@@ -54,7 +54,7 @@ class CLIMainTestCase(unittest.TestCase):
         service_names = ['web', 'db']
         timeout = 12
 
-        with mock.patch('compose.cli.main.signal', autospec=True) as mock_signal:
+        with mock.patch('compose.cli.main.signals.signal', autospec=True) as mock_signal:
             attach_to_logs(project, log_printer, service_names, timeout)
 
         assert mock_signal.signal.mock_calls == [