diff --git a/test/unit/__init__.py b/test/unit/__init__.py index 7afc5a946e..e367737084 100644 --- a/test/unit/__init__.py +++ b/test/unit/__init__.py @@ -32,6 +32,7 @@ import eventlet from eventlet.green import socket from tempfile import mkdtemp from shutil import rmtree +import signal import json @@ -1061,3 +1062,20 @@ def mocked_http_conn(*args, **kwargs): def make_timestamp_iter(): return iter(Timestamp(t) for t in itertools.count(int(time.time()))) + + +class Timeout(object): + def __init__(self, seconds): + self.seconds = seconds + + def __enter__(self): + signal.signal(signal.SIGALRM, self._exit) + signal.alarm(self.seconds) + + def __exit__(self, type, value, traceback): + signal.alarm(0) + + def _exit(self, signum, frame): + class TimeoutException(Exception): + pass + raise TimeoutException diff --git a/test/unit/cli/test_ringbuilder.py b/test/unit/cli/test_ringbuilder.py index 1f2b9494a5..b61907c967 100644 --- a/test/unit/cli/test_ringbuilder.py +++ b/test/unit/cli/test_ringbuilder.py @@ -29,6 +29,8 @@ from swift.cli.ringbuilder import EXIT_SUCCESS, EXIT_WARNING, EXIT_ERROR from swift.common import exceptions from swift.common.ring import RingBuilder +from test.unit import Timeout + class RunSwiftRingBuilderMixin(object): @@ -1955,6 +1957,29 @@ class TestCommands(unittest.TestCase, RunSwiftRingBuilderMixin): argv = ["-safe", self.tmpfile] self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv) + def test_remove_all_devices(self): + # Would block without the 'yes' argument + self.create_sample_ring() + argv = ["", self.tmpfile, "remove", "--weight", "100", "--yes"] + with Timeout(5): + self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv) + + def test_set_info_all_devices(self): + # Would block without the 'yes' argument + self.create_sample_ring() + argv = ["", self.tmpfile, "set_info", "--weight", "100", + "--change-meta", "something", "--yes"] + with Timeout(5): + self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv) + + def test_set_weight_all_devices(self): + # Would block without the 'yes' argument + self.create_sample_ring() + argv = ["", self.tmpfile, "set_weight", + "--weight", "100", "200", "--yes"] + with Timeout(5): + self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv) + class TestRebalanceCommand(unittest.TestCase, RunSwiftRingBuilderMixin):