Wait for all containers to exit when running 'up' interactively

Signed-off-by: Aanand Prasad <[email protected]>

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:

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())),
+        )