diff --git a/compose/cli/multiplexer.py b/compose/cli/multiplexer.py index ab7482e1d..34b55133c 100644 --- a/compose/cli/multiplexer.py +++ b/compose/cli/multiplexer.py @@ -7,8 +7,6 @@ except ImportError: from queue import Queue, Empty # Python 3.x -# Yield STOP from an input generator to stop the -# top-level loop without processing any more input. STOP = object() @@ -20,16 +18,17 @@ class Multiplexer(object): def __init__(self, iterators): self.iterators = iterators + self._num_running = len(iterators) self.queue = Queue() def loop(self): self._init_readers() - while True: + while self._num_running > 0: try: item = self.queue.get(timeout=0.1) if item is STOP: - break + self._num_running -= 1 else: yield item except Empty: diff --git a/tests/unit/multiplexer_test.py b/tests/unit/multiplexer_test.py new file mode 100644 index 000000000..100b8f0c2 --- /dev/null +++ b/tests/unit/multiplexer_test.py @@ -0,0 +1,28 @@ +import unittest + +from compose.cli.multiplexer import Multiplexer + + +class MultiplexerTest(unittest.TestCase): + def test_no_iterators(self): + mux = Multiplexer([]) + self.assertEqual([], list(mux.loop())) + + def test_empty_iterators(self): + mux = Multiplexer([ + (x for x in []), + (x for x in []), + ]) + + self.assertEqual([], list(mux.loop())) + + def test_aggregates_output(self): + mux = Multiplexer([ + (x for x in [0, 2, 4]), + (x for x in [1, 3, 5]), + ]) + + self.assertEqual( + [0, 1, 2, 3, 4, 5], + sorted(list(mux.loop())), + )