Add progress_callback to create.

blakerouse · blakerouse · commit 41377b9aa0c1 · 2016-10-26T13:08:50.000-04:00
diff --git a/maas/client/viscera/boot_resources.py b/maas/client/viscera/boot_resources.py
@@ -9,6 +9,7 @@
 import hashlib
 import httplib2
 import io
+from typing import Callable
 
 from . import (
     check,
@@ -117,7 +118,7 @@ def create(
             cls, name: str, architecture: str, content: io.IOBase, *,
             title: str=None,
             filetype: BootResourceFiletype=BootResourceFiletype.TGZ,
-            chunk_size=(1 << 22)):
+            chunk_size=(1 << 22), progress_callback: Callable=None):
         """Create a `BootResource`.
 
         Creates an uploaded boot resource with `content`. The `content` is
@@ -139,6 +140,11 @@ def create(
         :param chunk_size: Size in bytes to upload to MAAS in chunks.
             (Default is 4 MiB).
         :type chunk_size: `int`
+        :param progress_callback: Called to inform the current progress of the
+            upload. One argument is passed with the progress as a precentage.
+            If the resource was already complete and no content
+            needed to be uploaded then this callback will never be called.
+        :type progress_callback: Callable
         :returns: Create boot resource.
         :rtype: `BootResource`.
         """
@@ -178,6 +184,11 @@ def create(
         elif chunk_size <= 0:
             raise ValueError(
                 "chunk_size must be greater than 0, not %d" % chunk_size)
+        if (progress_callback is not None and
+                not isinstance(progress_callback, Callable)):
+            raise TypeError(
+                "progress_callback must be a Callable, not %s" % (
+                type(progress_callback).__name__))
 
         size, sha256 = calc_size_and_sha265(content, chunk_size)
         resource = cls._object(cls._handler.create(
@@ -191,12 +202,13 @@ def create(
             return resource
         else:
             # Upload in chunks and reload boot resource.
-            cls._upload_chunks(rfile, content, chunk_size)
+            cls._upload_chunks(rfile, content, chunk_size, progress_callback)
             return cls._object.read(resource.id)
 
     @classmethod
     def _upload_chunks(
-            cls, rfile: BootResourceFile, content: io.IOBase, chunk_size: int):
+            cls, rfile: BootResourceFile, content: io.IOBase, chunk_size: int,
+            progress_callback: Callable=None):
         """Upload the `content` to `rfile` in chunks using `chunk_size`."""
         content.seek(0, io.SEEK_SET)
         upload_uri = "%s%s" % (
@@ -207,6 +219,8 @@ def _upload_chunks(
             length = len(buf)
             if length > 0:
                 cls._put_chunk(upload_uri, buf)
+                if progress_callback is not None:
+                    progress_callback(length / rfile.size)
             if length != chunk_size:
                 break
 
diff --git a/maas/client/viscera/tests/test_boot_resources.py b/maas/client/viscera/tests/test_boot_resources.py
@@ -276,6 +276,16 @@ def test__create_raises_ValueError_when_chunk_size_is_less_than_zero(self):
         self.assertEquals(
             "chunk_size must be greater than 0, not -1", str(error))
 
+    def test__create_raises_TypeError_when_progress_callback_no_callable(self):
+        BootResources = make_origin().BootResources
+
+        buf = io.BytesIO(b"")
+        error = self.assertRaises(
+            TypeError, BootResources.create,
+            "os/release", "arch/subarch", buf, progress_callback='')
+        self.assertEquals(
+            "progress_callback must be a Callable, not str", str(error))
+
     def test__create_calls_create_on_handler_does_nothing_if_complete(self):
         resource_id = random.randint(0, 100)
         name = "%s/%s" % (
@@ -418,9 +428,14 @@ def test__create_uploads_in_chunks_and_reloads_resource(self):
             httplib2.Response({"status": 200}), b"")
         self.patch(boot_resources.httplib2, "Http", http)
 
+        # Progress handler called on each chunk.
+        progress_handler = MagicMock()
+
+        # Create and upload the resource.
         resource = BootResources.create(
             name, architecture, buf,
-            title=title, filetype=filetype, chunk_size=chunk_size)
+            title=title, filetype=filetype, chunk_size=chunk_size,
+            progress_callback=progress_handler)
 
         # Check that returned resource is correct and updated.
         self.assertThat(resource, MatchesStructure.byEquality(
@@ -445,6 +460,13 @@ def test__create_uploads_in_chunks_and_reloads_resource(self):
         ]
         self.assertEquals(calls, http.return_value.request.call_args_list)
 
+        # Check that progress handler was called on each chunk.
+        calls = [
+            call(len(data[0 + i:chunk_size + i]) / len(data))
+            for i in range(0, len(data), chunk_size)
+        ]
+        self.assertEquals(calls, progress_handler.call_args_list)
+
     def test__create_raises_CallError_on_chunk_upload_failure(self):
         resource_id = random.randint(0, 100)
         name = "%s/%s" % (
