Introduce an interface for progress bars.

This makes typing a bit more obvious.

Bug: T212783
Change-Id: Ibadbe64c9c8de19c572b53468b4fdb5fd1ea8228

cumin/tests/unit/transports/test_clustershell.py

@@ -5,7 +5,7 @@
 import pytest
 
 from cumin import CuminError, nodeset
-from cumin.transports import BaseWorker, clustershell, Command, ProgressBars, State, Target, WorkerError
+from cumin.transports import BaseWorker, clustershell, Command, State, Target, TqdmProgressBars, WorkerError
 
 
 def test_node_class_instantiation():
@@ -202,7 +202,7 @@ class ConcreteBaseEventHandler(clustershell.BaseEventHandler):
     def __init__(self, nodes, commands, **kwargs):
         """Initialize progress bars."""
         super().__init__(nodes, commands, **kwargs)
-        self.progress = mock.Mock(spec_set=ProgressBars)
+        self.progress = mock.Mock(spec_set=TqdmProgressBars)
 
     def close(self, task):
         """Required by the BaseEventHandler class."""
@@ -284,7 +284,7 @@ def setup_method(self, _, tqdm, logger):  # pylint: disable=arguments-differ
         """Initialize default properties and instances."""
         super().setup_method()
         self.handler = clustershell.SyncEventHandler(self.target, self.commands, success_threshold=1)
-        self.handler.progress = mock.Mock(spec_set=ProgressBars)
+        self.handler.progress = mock.Mock(spec_set=TqdmProgressBars)
         self.worker.eh = self.handler
         self.logger = logger
         assert not tqdm.write.called
@@ -370,7 +370,7 @@ def test_close(self, tqdm):  # pylint: disable=arguments-differ
 class TestAsyncEventHandler(TestBaseEventHandler):
     """AsyncEventHandler test class."""
 
-    @mock.patch('cumin.transports.clustershell.ProgressBars')
+    @mock.patch('cumin.transports.clustershell.TqdmProgressBars')
     @mock.patch('cumin.transports.clustershell.logging')
     @mock.patch('cumin.transports.clustershell.tqdm')
     def setup_method(self, _, tqdm, logger, progress):  # pylint: disable=arguments-differ,unused-argument

cumin/tests/unit/transports/test_init.py

@@ -9,7 +9,7 @@
 import cumin  # noqa: F401 (dynamically used in TestCommand)
 
 from cumin import transports
-from cumin.transports import ProgressBars
+from cumin.transports import TqdmProgressBars
 
 
 class ConcreteBaseWorker(transports.BaseWorker):
@@ -546,37 +546,37 @@ class TestProgressBars:
 
     def test_init_intialize_progress_bars_with_correct_size(self, tqdm):
         """Progress bars are initialized at the correct size."""
-        progress = ProgressBars()
+        progress = TqdmProgressBars()
         progress.init(10)
 
         assert tqdm.call_count == 2
         _, kwargs = tqdm.call_args
         assert kwargs['total'] == 10
 
-    def test_progress_bars_are_closed(self, tqdm):  # pylint: disable=unused-argument
+    def test_progress_bars_are_closed(self, tqdm):
         """Progress bars are closed."""
-        progress = ProgressBars()
+        progress = TqdmProgressBars()
         progress.init(10)
 
         progress.close()
 
-        assert progress.pbar_ok.close.called  # pylint: disable=no-member
-        assert progress.pbar_ko.close.called  # pylint: disable=no-member
+        assert tqdm.mock_calls[-2] == mock.call().close()
+        assert tqdm.mock_calls[-1] == mock.call().close()
 
-    def test_progress_bar_is_updated_on_success(self, tqdm):  # pylint: disable=unused-argument
+    def test_progress_bar_is_updated_on_success(self, tqdm):
         """Progress bar is updated on success."""
-        progress = ProgressBars()
+        progress = TqdmProgressBars()
         progress.init(10)
 
         progress.update_success(5)
 
-        assert progress.pbar_ok.update.called_once_with(5)  # pylint: disable=no-member
+        assert mock.call().update(5) in tqdm.mock_calls
 
-    def test_progress_bar_is_updated_on_failure(self, tqdm):  # pylint: disable=unused-argument
+    def test_progress_bar_is_updated_on_failure(self, tqdm):
         """Progress bar is updated on failure."""
-        progress = ProgressBars()
+        progress = TqdmProgressBars()
         progress.init(10)
 
         progress.update_failed(3)
 
-        assert progress.pbar_ko.update.called_once_with(3)  # pylint: disable=no-member
+        assert tqdm.mock_calls[-1] == mock.call().update(3)

cumin/transports/__init__.py

@@ -5,6 +5,7 @@
 import sys
 
 from abc import ABCMeta, abstractmethod
+from typing import Callable, Optional
 
 from ClusterShell.NodeSet import NodeSet
 from tqdm import tqdm
@@ -750,75 +751,125 @@ def raise_error(property_name, message, value):
         property_name=property_name, message=message, value_type=type(value), value=value))
 
 
-class ProgressBars:
-    """Progress bars for the status of successful / failed hosts.
+class BaseExecutionProgress(metaclass=ABCMeta):
+    """Listener interface to consume notification of the status of successful / failed hosts.
 
-    The ProgressBars needs to be notified of the total number of hosts when the
+    The listener needs to be notified of the total number of hosts when the
     operation starts, and then notified of successes and failures.
+
     """
 
-    def __init__(self):
+    @abstractmethod
+    def init(self, num_hosts: int) -> None:
+        """Initialize the progress bars.
+
+        Arguments:
+            num_hosts (int): the total number of hosts
+
+        """
+
+    @abstractmethod
+    def close(self) -> None:
+        """Closes the progress bars."""
+
+    @abstractmethod
+    def update_success(self, num_hosts: int = 1) -> None:
+        """Updates the number of successful hosts.
+
+        Arguments:
+            num_hosts (int): increment to the number of hosts that have completed successfully
+
+        """
+
+    @abstractmethod
+    def update_failed(self, num_hosts: int = 1) -> None:
+        """Updates the number of failed hosts.
+
+        Arguments:
+            num_hosts (int): increment to the number of hosts that have completed in error
+
+        """
+
+
+class TqdmProgressBars(BaseExecutionProgress):
+    """Progress bars based on TQDM."""
+
+    def __init__(self) -> None:
         """Create the progress bars.
 
         Note:
             the progress bars themselves are not initalized at object creation. ``init()`` needs to be called before
             using the progress bars.
 
         """
-        self.pbar_ok = None
-        self.pbar_ko = None
-        self.bar_format = ('{desc} |{bar}| {percentage:3.0f}% ({n_fmt}/{total_fmt}) '
-                           '[{elapsed}<{remaining}, {rate_fmt}]')
+        self._pbar_success: Optional[tqdm] = None
+        self._pbar_failed: Optional[tqdm] = None
+        self._bar_format = ('{desc} |{bar}| {percentage:3.0f}% ({n_fmt}/{total_fmt}) '
+                            '[{elapsed}<{remaining}, {rate_fmt}]')
 
-    def init(self, num_hosts):
+    def init(self, num_hosts: int) -> None:
         """Initialize the progress bars.
 
         Arguments:
             num_hosts (int): the total number of hosts
 
         """
-        self.pbar_ok = tqdm(desc='PASS', total=num_hosts, leave=True, unit='hosts', dynamic_ncols=True,
-                            bar_format=Colored.green(self.bar_format), file=sys.stderr)
-        self.pbar_ok.refresh()
-        self.pbar_ko = tqdm(desc='FAIL', total=num_hosts, leave=True, unit='hosts', dynamic_ncols=True,
-                            bar_format=Colored.red(self.bar_format), file=sys.stderr)
-        self.pbar_ko.refresh()
+        self._pbar_success = self._tqdm(num_hosts, 'PASS', Colored.green)
+        self._pbar_failed = self._tqdm(num_hosts, 'FAIL', Colored.red)
+
+    def _tqdm(self, num_hosts: int, desc: str, color: Callable[[str], str]) -> tqdm:
+        pbar = tqdm(desc=desc, total=num_hosts, leave=True, unit='hosts', dynamic_ncols=True,
+                    bar_format=color(self._bar_format), file=sys.stderr)
+        pbar.refresh()
+        return pbar
 
-    def close(self):
+    def close(self) -> None:
         """Closes the progress bars."""
-        self.pbar_ok.close()
-        self.pbar_ko.close()
+        self._success.close()
+        self._failed.close()
 
-    def update_success(self, num_hosts=1):
+    def update_success(self, num_hosts: int = 1) -> None:
         """Updates the number of successful hosts.
 
         Arguments:
             num_hosts (int): increment to the number of hosts that have completed successfully
 
         """
-        self.pbar_ok.update(num_hosts)
+        self._success.update(num_hosts)
 
-    def update_failed(self, num_hosts=1):
+    def update_failed(self, num_hosts: int = 1) -> None:
         """Updates the number of failed hosts.
 
         Arguments:
             num_hosts (int): increment to the number of hosts that have completed in error
 
         """
-        self.pbar_ko.update(num_hosts)
+        self._failed.update(num_hosts)
+
+    @property
+    def _success(self) -> tqdm:
+        if self._pbar_success is None:
+            raise ValueError('init() should be called before any other operation')
+        return self._pbar_success
+
+    @property
+    def _failed(self) -> tqdm:
+        if self._pbar_failed is None:
+            raise ValueError('init() should be called before any other operation')
+        return self._pbar_failed
 
 
-class NoProgress:
-    """Used as a null object to disable the display of progress bars."""
+class NoProgress(BaseExecutionProgress):
+    """Used as a null object to disable the display of execution progress."""
 
-    def init(self, num_hosts):
+    def init(self, num_hosts: int) -> None:
         """Does nothing."""
 
-    def close(self):
+    def close(self) -> None:
         """Does nothing."""
 
-    def update_success(self, num_hosts=1):
+    def update_success(self, num_hosts: int = 1) -> None:
         """Does nothing."""
 
-    def update_failed(self, num_hosts=1):
+    def update_failed(self, num_hosts: int = 1) -> None:
         """Does nothing."""

cumin/transports/clustershell.py

@@ -11,7 +11,8 @@
 
 from cumin import nodeset, nodeset_fromlist
 from cumin.color import Colored
-from cumin.transports import BaseWorker, Command, NoProgress, ProgressBars, raise_error, State, Target, WorkerError
+from cumin.transports import (BaseExecutionProgress, BaseWorker, Command, NoProgress, raise_error, State,
+                              Target, TqdmProgressBars, WorkerError)
 
 
 class ClusterShellWorker(BaseWorker):
@@ -405,8 +406,7 @@ def __init__(self, target: Target, commands: List[Command], success_threshold: f
         for node_name in target.first_batch:
             self.nodes[node_name].state.update(State.scheduled)
 
-        # TODO: introduce a super type to ProgressBars / NoProgress
-        self.progress: Union[ProgressBars, NoProgress] = ProgressBars() if progress_bars else NoProgress()
+        self.progress: BaseExecutionProgress = TqdmProgressBars() if progress_bars else NoProgress()
         self.reporter = Reporter()
 
     def close(self, task):