Merge pull request #514 from coreos-inc/python-registry-v2-patch-tests

Add PATCH tests for resumable upload

test/registry_tests.py

@@ -1,7 +1,9 @@
 import unittest
 import requests
 import os
-
+import math
+import random
+import string
 
 import Crypto.Random
 from flask import request, jsonify
@@ -219,38 +221,41 @@ class V1RegistryPushMixin(V1RegistryMixin):
                  data=json.dumps(data), auth=auth,
                  expected_code=201)
 
-    for image in images:
+    last_image_id = None
+    for image_id, _ in images.iteritems():
+      last_image_id = image_id
+
       # PUT /v1/images/{imageID}/json
-      self.conduct('PUT', '/v1/images/%s/json' % image['id'],
+      self.conduct('PUT', '/v1/images/%s/json' % image_id,
-                   data=json.dumps(image), auth='sig')
+                   data=json.dumps({'id': image_id}), auth='sig')
 
       # PUT /v1/images/{imageID}/layer
       tar_file_info = tarfile.TarInfo(name='image_name')
       tar_file_info.type = tarfile.REGTYPE
-      tar_file_info.size = len(image['id'])
+      tar_file_info.size = len(image_id)
 
       layer_data = StringIO()
 
       tar_file = tarfile.open(fileobj=layer_data, mode='w|gz')
-      tar_file.addfile(tar_file_info, StringIO(image['id']))
+      tar_file.addfile(tar_file_info, StringIO(image_id))
       tar_file.close()
 
       layer_bytes = layer_data.getvalue()
       layer_data.close()
 
-      self.conduct('PUT', '/v1/images/%s/layer' % image['id'],
+      self.conduct('PUT', '/v1/images/%s/layer' % image_id,
                    data=StringIO(layer_bytes), auth='sig')
 
       # PUT /v1/images/{imageID}/checksum
-      checksum = compute_simple(StringIO(layer_bytes), json.dumps(image))
+      checksum = compute_simple(StringIO(layer_bytes), json.dumps({'id': image_id}))
-      self.conduct('PUT', '/v1/images/%s/checksum' % image['id'],
+      self.conduct('PUT', '/v1/images/%s/checksum' % image_id,
                    headers={'X-Docker-Checksum-Payload': checksum},
                    auth='sig')
 
 
     # PUT /v1/repositories/{namespace}/{repository}/tags/latest
     self.conduct('PUT', '/v1/repositories/%s/%s/tags/latest' % (namespace, repository),
-                 data='"' + images[0]['id'] + '"',
+                 data='"' + last_image_id + '"',
                  auth='sig')
 
     # PUT /v1/repositories/{namespace}/{repository}/images
@@ -318,18 +323,28 @@ class V2RegistryPushMixin(V2RegistryMixin):
     self.do_auth(username, password, namespace, repository, scopes=['push', 'pull'])
 
     # Build a fake manifest.
-    images = [('somelayer', 'some fake data')]
-
     tag_name = 'latest'
     builder = SignedManifestBuilder(namespace, repository, tag_name)
-    for image_id, contents in images:
+    for image_id, contents in images.iteritems():
-      checksum = 'sha256:' + hashlib.sha256(contents).hexdigest()
+      if isinstance(contents, dict):
+        full_contents = contents['contents']
+      else:
+        full_contents = contents
+
+      checksum = 'sha256:' + hashlib.sha256(full_contents).hexdigest()
       builder.add_layer(checksum, json.dumps({'id': image_id, 'data': contents}))
 
     # Push the image's layers.
-    for image_id, contents in images:
+    for image_id, contents in images.iteritems():
+      chunks = None
+      if isinstance(contents, dict):
+        full_contents = contents['contents']
+        chunks = contents['chunks']
+      else:
+        full_contents = contents
+
       # Layer data should not yet exist.
-      checksum = 'sha256:' + hashlib.sha256(contents).hexdigest()
+      checksum = 'sha256:' + hashlib.sha256(full_contents).hexdigest()
       self.conduct('HEAD', '/v2/%s/%s/blobs/%s' % (namespace, repository, checksum),
                    expected_code=404, auth='jwt')
 
@@ -340,7 +355,14 @@ class V2RegistryPushMixin(V2RegistryMixin):
       location = response.headers['Location'][len(self.get_server_url()):]
 
       # PATCH the image data into the layer.
+      if chunks is None:
         self.conduct('PATCH', location, data=contents, expected_code=204, auth='jwt')
+      else:
+        for chunk in chunks:
+          (start_byte, end_byte) = chunk
+          contents_chunk = full_contents[start_byte:end_byte]
+          self.conduct('PATCH', location, data=contents_chunk, expected_code=204, auth='jwt',
+                       headers={'Range': 'bytes=%s-%s' % (start_byte, end_byte)})
 
       # Finish the layer upload with a PUT.
       self.conduct('PUT', location, params=dict(digest=checksum), expected_code=201, auth='jwt')
@@ -371,18 +393,25 @@ class V2RegistryPullMixin(V2RegistryMixin):
     response = self.conduct('GET', '/v2/%s/%s/manifests/%s' % (namespace, repository, tag_name),
                             auth='jwt')
     manifest_data = json.loads(response.text)
+    blobs = {}
+
     for layer in manifest_data['fsLayers']:
       blob_id = layer['blobSum']
-      self.conduct('GET', '/v2/%s/%s/blobs/%s' % (namespace, repository, blob_id),
+      result = self.conduct('GET', '/v2/%s/%s/blobs/%s' % (namespace, repository, blob_id),
                             expected_code=200, auth='jwt')
 
+      blobs[blob_id] = result.text
+
+    return blobs
+
 
 class RegistryTestsMixin(object):
   def test_pull_publicrepo_anonymous(self):
     # Add a new repository under the public user, so we have a real repository to pull.
-    images = [{
+    images = {
-      'id': 'onlyimagehere'
+      'someid': 'onlyimagehere'
-    }]
+    }
+
     self.do_push('public', 'newrepo', 'public', 'password', images)
     self.clearSession()
 
@@ -401,9 +430,10 @@ class RegistryTestsMixin(object):
 
   def test_pull_publicrepo_devtable(self):
     # Add a new repository under the public user, so we have a real repository to pull.
-    images = [{
+    images = {
-      'id': 'onlyimagehere'
+      'someid': 'onlyimagehere'
-    }]
+    }
+
     self.do_push('public', 'newrepo', 'public', 'password', images)
     self.clearSession()
 
@@ -422,9 +452,10 @@ class RegistryTestsMixin(object):
 
   def test_pull_private_repo(self):
     # Add a new repository under the devtable user, so we have a real repository to pull.
-    images = [{
+    images = {
-      'id': 'onlyimagehere'
+      'someid': 'onlyimagehere'
-    }]
+    }
+
     self.do_push('devtable', 'newrepo', 'devtable', 'password', images)
     self.clearSession()
 
@@ -441,9 +472,10 @@ class RegistryTestsMixin(object):
     # Turn off anonymous access.
     with TestFeature(self, 'ANONYMOUS_ACCESS', False):
       # Add a new repository under the public user, so we have a real repository to pull.
-      images = [{
+      images = {
-        'id': 'onlyimagehere'
+        'someid': 'onlyimagehere'
-      }]
+      }
+
       self.do_push('public', 'newrepo', 'public', 'password', images)
       self.clearSession()
 
@@ -464,9 +496,10 @@ class RegistryTestsMixin(object):
     # Turn off anonymous access.
     with TestFeature(self, 'ANONYMOUS_ACCESS', False):
       # Add a new repository under the public user, so we have a real repository to pull.
-      images = [{
+      images = {
-        'id': 'onlyimagehere'
+        'someid': 'onlyimagehere'
-      }]
+      }
+
       self.do_push('public', 'newrepo', 'public', 'password', images)
       self.clearSession()
 
@@ -482,9 +515,10 @@ class RegistryTestsMixin(object):
     # Turn off anonymous access.
     with TestFeature(self, 'ANONYMOUS_ACCESS', False):
       # Add a new repository under the public user, so we have a real repository to pull.
-      images = [{
+      images = {
-        'id': 'onlyimagehere'
+        'someid': 'onlyimagehere'
-      }]
+      }
+
       self.do_push('public', 'newrepo', 'public', 'password', images)
       self.clearSession()
 
@@ -509,9 +543,10 @@ class RegistryTestsMixin(object):
 
 
   def test_create_repo_creator_user(self):
-    images = [{
+    images = {
-      'id': 'onlyimagehere'
+      'someid': 'onlyimagehere'
-    }]
+    }
+
     self.do_push('buynlarge', 'newrepo', 'creator', 'password', images)
 
     # Pull the repository as devtable, which should succeed because the repository is owned by the
@@ -525,9 +560,10 @@ class RegistryTestsMixin(object):
     resp = self.conduct('GET', '/api/v1/organization/buynlarge/robots/ownerbot')
     robot_token = json.loads(resp.text)['token']
 
-    images = [{
+    images = {
-      'id': 'onlyimagehere'
+      'someid': 'onlyimagehere'
-    }]
+    }
+
     self.do_push('buynlarge', 'newrepo', 'buynlarge+ownerbot', robot_token, images)
 
     # Pull the repository as devtable, which should succeed because the repository is owned by the
@@ -541,9 +577,10 @@ class RegistryTestsMixin(object):
     resp = self.conduct('GET', '/api/v1/organization/buynlarge/robots/creatorbot')
     robot_token = json.loads(resp.text)['token']
 
-    images = [{
+    images = {
-      'id': 'onlyimagehere'
+      'someid': 'onlyimagehere'
-    }]
+    }
+
     self.do_push('buynlarge', 'newrepo', 'buynlarge+creatorbot', robot_token, images)
 
     # Pull the repository as devtable, which should succeed because the repository is owned by the
@@ -559,6 +596,51 @@ class V2RegistryTests(V2RegistryPullMixin, V2RegistryPushMixin, RegistryTestsMix
                       RegistryTestCaseMixin, LiveServerTestCase):
   """ Tests for V2 registry. """
 
+  def test_partial_upload_below_5mb(self):
+    chunksize = 1024 * 1024 * 2
+    size = chunksize * 3
+    contents = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(size))
+
+    chunk_count = int(math.ceil((len(contents) * 1.0) / chunksize))
+    chunks = [(index * chunksize, (index + 1)*chunksize) for index in range(chunk_count)]
+
+    images = {
+      'someid': {
+        'contents': contents,
+        'chunks': chunks
+      }
+    }
+
+    # Push the chunked upload.
+    self.do_push('devtable', 'newrepo', 'devtable', 'password', images)
+
+    # Pull the image back and verify the contents.
+    blobs = self.do_pull('devtable', 'newrepo', 'devtable', 'password')
+    self.assertEquals(len(blobs.items()), 1)
+    self.assertEquals(blobs.items()[0][1], contents)
+
+  def test_partial_upload_resend_below_5mb(self):
+    size = 1024 * 1024 * 2
+    contents = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(size))
+
+    chunks = [(0, 10), (0, 100), (100, size)]
+
+    images = {
+      'someid': {
+        'contents': contents,
+        'chunks': chunks
+      }
+    }
+
+    # Push the chunked upload.
+    self.do_push('devtable', 'newrepo', 'devtable', 'password', images)
+
+    # Pull the image back and verify the contents.
+    blobs = self.do_pull('devtable', 'newrepo', 'devtable', 'password')
+    self.assertEquals(len(blobs.items()), 1)
+    self.assertEquals(blobs.items()[0][1], contents)
+
+
 class V1PushV2PullRegistryTests(V2RegistryPullMixin, V1RegistryPushMixin, RegistryTestsMixin,
                                 RegistryTestCaseMixin, LiveServerTestCase):
   """ Tests for V1 push, V2 pull registry. """