Merge lp:~barry/ubuntu-system-image/lp1423622 into lp:~registry/ubuntu-system-image/client

[Barry Warsaw]

--progress

[Barry Warsaw (barry)]

I'll add NEWS and update the manpage when I merge this into trunk.

[Michael Vogt (mvo)]

Thanks for this work Barry, this looks very nice!

There are two changes I would love to see:
1. progress should be unbuffered (that was wrong in my initial branch, sorry). I.e. instead of print(), could you please use sys.stdout.write() and sys.stdout.flush() ?

2. if a exception happens in main.py (cf. http://bazaar.launchpad.net/~ubuntu-system-image/ubuntu-system-image/client/view/head:/systemimage/main.py#L381) it would be nice to get a json error on stdout. Not sure how well this fits into the new schema, I guess it is ok if I get a non-json error on stderr that I can just display to the user.

[Barry Warsaw (barry)]

On Mar 05, 2015, at 08:05 AM, Michael Vogt wrote:

>1. progress should be unbuffered (that was wrong in my initial branch,
>sorry). I.e. instead of print(), could you please use sys.stdout.write() and
>sys.stdout.flush() ?

Good point, I should have noticed that. I'll fix that asap.

>2. if a exception happens in main.py
>(cf. http://bazaar.launchpad.net/~ubuntu-system-image/ubuntu-system-image/client/view/head:/systemimage/main.py#L381)
>it would be nice to get a json error on stdout. Not sure how well this fits
>into the new schema, I guess it is ok if I get a non-json error on stderr
>that I can just display to the user.

Okay, we chatted on IRC and decided we don't need much here. I've been
meaning to add a simple message to stderr in this case anyway, so that's what
I'll do now.

=== modified file 'systemimage/api.py'
--- systemimage/api.py	2015-02-09 23:28:14 +0000
+++ systemimage/api.py	2015-03-04 20:03:16 +0000
@@ -26,7 +26,6 @@
 
 from systemimage.apply import factory_reset, production_reset
 from systemimage.state import State
-from unittest.mock import patch
 
 
 log = logging.getLogger('systemimage')
@@ -119,8 +118,12 @@
     def download(self):
         """Download the available update."""
         # We only want callback progress during the actual download.
-        with patch.object(self._state.downloader, 'callback', self._callback):
+        old_callbacks = self._state.downloader.callbacks[:]
+        try:
+            self._state.downloader.callbacks = [self._callback]
             self._state.run_until('apply')
+        finally:
+            self._state.downloader.callbacks = old_callbacks
 
     def apply(self):
         """Apply the update."""

=== modified file 'systemimage/curl.py'
--- systemimage/curl.py	2015-02-13 21:44:50 +0000
+++ systemimage/curl.py	2015-03-04 20:03:16 +0000
@@ -130,7 +130,9 @@
     """The PyCURL based download manager."""
 
     def __init__(self, callback=None):
-        super().__init__(callback)
+        super().__init__()
+        if callback is not None:
+            self.callbacks.append(callback)
         self._pausables = []
         self._paused = False
 

=== modified file 'systemimage/download.py'
--- systemimage/download.py	2015-02-13 21:44:50 +0000
+++ systemimage/download.py	2015-03-04 20:03:16 +0000
@@ -71,7 +71,7 @@
 class DownloadManagerBase:
     """Base class for all download managers."""
 
-    def __init__(self, callback=None):
+    def __init__(self):
         """
         :param callback: If given, a function that is called every so often
             during downloading.
@@ -79,10 +79,14 @@
             of bytes received so far, and the total amount of bytes to be
             downloaded.
         """
-        self._queued_cancel = False
-        self.callback = callback
+        # This is a list of functions that are called every so often during
+        # downloading.  Functions in this list take two arguments, the number
+        # of bytes received so far, and the total amount of bytes to be
+        # downloaded.
+        self.callbacks = []
         self.total = 0
         self.received = 0
+        self._queued_cancel = False
 
     def __repr__(self): # pragma: no cover
         return '<{} at 0x{:x}>'.format(self.__class__.__name__, id(self))
@@ -128,9 +132,9 @@
     def _do_callback(self):
         # Be defensive, so yes, use a bare except.  If an exception occurs in
         # the callback, log it, but continue onward.
-        if self.callback is not None:
+        for callback in self.callbacks:
             try:
-                self.callback(self.received, self.total)
+                callback(self.received, self.total)
             except:
                 log.exception('Exception in progress callback')
 

=== modified file 'systemimage/main.py'
--- systemimage/main.py	2015-02-09 18:46:24 +0000
+++ systemimage/main.py	2015-03-04 20:03:16 +0000
@@ -22,6 +22,7 @@
 
 
 import sys
+import json
 import logging
 import argparse
 
@@ -43,6 +44,39 @@
 
 DEFAULT_CONFIG_D = '/etc/system-image/config.d'
 COLON = ':'
+LINE_LENGTH = 78
+
+
+class _DotsProgress:
+    def __init__(self):
+        self._dot_count = 0
+
+    def callback(self, received, total):
+        received = int(received)
+        total = int(total)
+        sys.stderr.write('.')
+        sys.stderr.flush()
+        self._dot_count += 1
+        show_dots = self._dot_count % LINE_LENGTH == 0
+        if show_dots or received >= total:          # pragma: no cover
+            sys.stderr.write('\n')
+            sys.stderr.flush()
+
+
+class _LogfileProgress:
+    def __init__(self, log):
+        self._log = log
+
+    def callback(self, received, total):
+        self._log.debug('received: {} of {} bytes', received, total)
+
+
+def _json_progress(received, total):
+    # For use with --progress=json output.  LP: #1423622
+    print(json.dumps(dict(
+        type='progress',
+        now=received,
+        total=total)))
 
 
 def main():
@@ -94,13 +128,18 @@
                         default=False, action='store_true',
                         help="""Calculate and print the upgrade path, but do
                                 not download or apply it""")
+    parser.add_argument('-v', '--verbose',
+                        default=0, action='count',
+                        help='Increase verbosity')
+    parser.add_argument('--progress',
+                        default=[], action='append',
+                        help="""Add a progress meter.  Available meters are:
+                                dots, logfile, and json.  Multiple --progress
+                                options are allowed.""")
     parser.add_argument('-p', '--percentage',
                         default=None, action='store',
                         help="""Override the device's phased percentage value
                                 during upgrade candidate calculation.""")
-    parser.add_argument('-v', '--verbose',
-                        default=0, action='count',
-                        help='Increase verbosity')
     parser.add_argument('--list-channels',
                         default=False, action='store_true',
                         help="""List all available channels, then exit""")
@@ -285,29 +324,20 @@
                 print('    {} (alias for: {})'.format(key, alias))
         return 0
 
-    # When verbosity is at 1, logging every progress signal from
-    # ubuntu-download-manager would be way too noisy.  OTOH, not printing
-    # anything leads some folks to think the process is just hung, since it
-    # can take a long time to download all the data files.  As a compromise,
-    # we'll output some dots to stderr at verbosity 1, but we won't log these
-    # dots since they would just be noise.  This doesn't have to be perfect.
-    if args.verbose == 1:                           # pragma: no cover
-        dot_count = 0
-        def callback(received, total):
-            nonlocal dot_count
-            sys.stderr.write('.')
-            sys.stderr.flush()
-            dot_count += 1
-            if dot_count % 78 == 0 or received >= total:
-                sys.stderr.write('\n')
-                sys.stderr.flush()
-    else:
-        def callback(received, total):
-            log.debug('received: {} of {} bytes', received, total)
-
     DBusGMainLoop(set_as_default=True)
     state = State(candidate_filter=candidate_filter)
-    state.downloader.callback = callback
+
+    for meter in args.progress:
+        if meter == 'dots':
+            state.downloader.callbacks.append(_DotsProgress().callback)
+        elif meter == 'json':
+            state.downloader.callbacks.append(_json_progress)
+        elif meter == 'logfile':
+            state.downloader.callbacks.append(_LogfileProgress(log).callback)
+        else:
+            parser.error('Unknown progress meter: {}'.format(meter))
+            assert 'parser.error() does not return' # pragma: no cover
+
     if args.dry_run:
         try:
             state.run_until('download_files')

=== modified file 'systemimage/tests/test_main.py'
--- systemimage/tests/test_main.py	2015-02-18 18:54:49 +0000
+++ systemimage/tests/test_main.py	2015-03-04 20:03:16 +0000
@@ -18,13 +18,14 @@
 __all__ = [
     'TestCLIDuplicateDestinations',
     'TestCLIFactoryReset',
-    'TestCLIProductionReset',
     'TestCLIFilters',
     'TestCLIListChannels',
     'TestCLIMain',
     'TestCLIMainDryRun',
     'TestCLIMainDryRunAliases',
     'TestCLINoReboot',
+    'TestCLIProductionReset',
+    'TestCLIProgress',
     'TestCLISettings',
     'TestCLISignatures',
     'TestDBusMain',
@@ -35,6 +36,7 @@
 import os
 import sys
 import dbus
+import json
 import stat
 import time
 import shutil
@@ -1480,3 +1482,127 @@
 WARNING: All GPG signature verifications have been disabled.
 Your upgrades are INSECURE.
 """)
+
+
+class TestCLIProgress(ServerTestBase):
+    INDEX_FILE = 'main.index_01.json'
+    CHANNEL_FILE = 'main.channels_01.json'
+    CHANNEL = 'stable'
+    DEVICE = 'nexus7'
+
+    @configuration
+    def test_dots_progress(self, config_d):
+        # --progress=dots prints a bunch of dots to stderr.
+        self._setup_server_keyrings()
+        stderr = StringIO()
+        with ExitStack() as resources:
+            resources.enter_context(
+                patch('systemimage.main.LINE_LENGTH', 10))
+            resources.enter_context(
+                patch('systemimage.main.sys.stderr', stderr))
+            resources.enter_context(
+                argv('-C', config_d, '-b', '0', '--no-reboot',
+                     '--progress', 'dots'))
+            cli_main()
+        # There should be some dots in the stderr.
+        self.assertGreater(stderr.getvalue().count('.'), 2)
+
+    @configuration
+    def test_json_progress(self, config_d):
+        # --progress=json prints some JSON to stdout.
+        self._setup_server_keyrings()
+        stdout = StringIO()
+        with ExitStack() as resources:
+            resources.enter_context(capture_print(stdout))
+            resources.enter_context(
+                argv('-C', config_d, '-b', '0', '--no-reboot',
+                     '--progress', 'json'))
+            cli_main()
+        # stdout is now filled with JSON goodness.  We can't assert too much
+        # about the contents though.
+        line_count = 0
+        for line in stdout.getvalue().splitlines():
+            line_count += 1
+            record = json.loads(line)
+            self.assertEqual(record['type'], 'progress')
+            self.assertIn('now', record)
+            self.assertIn('total', record)
+        self.assertGreater(line_count, 4)
+
+    @configuration
+    def test_logfile_progress(self, config_d):
+        # --progress=logfile dumps some messages to the log file.
+        self._setup_server_keyrings()
+        log_mock = MagicMock()
+        from systemimage.main import _LogfileProgress
+        class Testable(_LogfileProgress):
+            def __init__(self, log):
+                super().__init__(log)
+                self._log = log_mock
+        with ExitStack() as resources:
+            resources.enter_context(
+                patch('systemimage.main._LogfileProgress', Testable))
+            resources.enter_context(
+                argv('-C', config_d, '-b', '0', '--no-reboot',
+                     '--progress', 'logfile'))
+            cli_main()
+        self.assertGreater(log_mock.debug.call_count, 4)
+        positional, keyword = log_mock.debug.call_args
+        self.assertTrue(positional[0].startswith('received: '))
+
+    @configuration
+    def test_all_progress(self, config_d):
+        # We can have more than one --progress flag.
+        self._setup_server_keyrings()
+        stdout = StringIO()
+        stderr = StringIO()
+        log_mock = MagicMock()
+        from systemimage.main import _LogfileProgress
+        class Testable(_LogfileProgress):
+            def __init__(self, log):
+                super().__init__(log)
+                self._log = log_mock
+        with ExitStack() as resources:
+            resources.enter_context(
+                patch('systemimage.main.LINE_LENGTH', 10))
+            resources.enter_context(
+                patch('systemimage.main.sys.stderr', stderr))
+            resources.enter_context(capture_print(stdout))
+            resources.enter_context(
+                patch('systemimage.main._LogfileProgress', Testable))
+            resources.enter_context(
+                argv('-C', config_d, '-b', '0', '--no-reboot',
+                     '--progress', 'dots',
+                     '--progress', 'json',
+                     '--progress', 'logfile'))
+            cli_main()
+        self.assertGreater(stderr.getvalue().count('.'), 2)
+        line_count = 0
+        for line in stdout.getvalue().splitlines():
+            line_count += 1
+            record = json.loads(line)
+            self.assertEqual(record['type'], 'progress')
+            self.assertIn('now', record)
+            self.assertIn('total', record)
+        self.assertGreater(line_count, 4)
+        self.assertGreater(log_mock.debug.call_count, 4)
+        positional, keyword = log_mock.debug.call_args
+        self.assertTrue(positional[0].startswith('received: '))
+
+    @configuration
+    def test_bad_progress(self, config_d):
+        # An unknown progress type results in an error.
+        stderr = StringIO()
+        with ExitStack() as resources:
+            resources.enter_context(
+                patch('systemimage.main.sys.stderr', stderr))
+            resources.enter_context(
+                argv('-C', config_d, '-b', '0', '--no-reboot',
+                     '--progress', 'not-a-meter'))
+            with self.assertRaises(SystemExit) as cm:
+                cli_main()
+            exit_code = cm.exception.code
+        self.assertEqual(exit_code, 2)
+        self.assertEqual(
+            stderr.getvalue().splitlines()[-1],
+            'system-image-cli: error: Unknown progress meter: not-a-meter')

=== modified file 'systemimage/udm.py'
--- systemimage/udm.py	2015-02-13 21:44:50 +0000
+++ systemimage/udm.py	2015-03-04 20:03:16 +0000
@@ -122,7 +122,9 @@
     """Download via ubuntu-download-manager (UDM)."""
 
     def __init__(self, callback=None):
-        super().__init__(callback)
+        super().__init__()
+        if callback is not None:
+            self.callbacks.append(callback)
         self._iface = None
 
     def _get_files(self, records, pausable):