diff --git a/compose/cli/main.py b/compose/cli/main.py index d10b95823..82cf05c9f 100644 --- a/compose/cli/main.py +++ b/compose/cli/main.py @@ -5,7 +5,6 @@ from __future__ import unicode_literals import json import logging import re -import signal import sys from inspect import getdoc from operator import attrgetter @@ -14,6 +13,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 @@ -682,20 +682,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) @@ -710,25 +709,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 84fe7f15b..3a3c89b05 100644 --- a/tests/acceptance/cli_test.py +++ b/tests/acceptance/cli_test.py @@ -90,7 +90,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 ab2368669..f62b2bcc3 100644 --- a/tests/unit/cli/main_test.py +++ b/tests/unit/cli/main_test.py @@ -55,7 +55,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 == [