Set image size after updating/adding locations

Based on current implement, the image size won't be updated if user
update/add locations against a queued image. This fix will set the
image size based on given location.

Closes-Bug: #1261624

Change-Id: I265c6a92274ab2f94e005a1ab50e01d2f0c2143e

glance/store/__init__.py

@@ -426,6 +426,17 @@ def _check_image_location(context, store_api, location):
     store_api.check_location_metadata(location['metadata'])
 
 
+def _set_image_size(context, image, locations):
+    if not image.size:
+        for location in locations:
+            size_from_backend = glance.store.get_size_from_backend(
+                context, location['url'])
+            if size_from_backend:
+                # NOTE(flwang): This assumes all locations have the same size
+                image.size = size_from_backend
+                break
+
+
 class ImageFactoryProxy(glance.domain.proxy.ImageFactory):
     def __init__(self, factory, context, store_api):
         self.context = context
@@ -483,6 +494,9 @@ class StoreLocations(collections.MutableSequence):
             raise exception.DuplicateLocation(location=location['url'])
 
         self.value.insert(i, location)
+        _set_image_size(self.image_proxy.context,
+                        self.image_proxy,
+                        [location])
 
     def pop(self, i=-1):
         location = self.value.pop(i)
@@ -521,6 +535,9 @@ class StoreLocations(collections.MutableSequence):
         _check_image_location(self.image_proxy.context,
                               self.image_proxy.store_api, location)
         self.value.__setitem__(i, location)
+        _set_image_size(self.image_proxy.context,
+                        self.image_proxy,
+                        [location])
 
     def __delitem__(self, i):
         location = None
@@ -599,7 +616,7 @@ def _locations_proxy(target, attr):
 
                 if value.count(location) > 1:
                     raise exception.DuplicateLocation(location=location['url'])
-
+            _set_image_size(self.context, getattr(self, target), value)
             return setattr(getattr(self, target), attr, list(value))
 
     def del_attr(self):

glance/tests/functional/v2/test_images.py

@@ -16,6 +16,7 @@
 #    under the License.
 
 import json
+import os
 import uuid
 
 import requests
@@ -1706,6 +1707,42 @@ class TestImages(functional.FunctionalTest):
 
         self.stop_servers()
 
+    def test_update_locations(self):
+        # Create an image
+        path = self._url('/v2/images')
+        headers = self._headers({'content-type': 'application/json'})
+        data = json.dumps({'name': 'image-1', 'disk_format': 'aki',
+                           'container_format': 'aki'})
+        response = requests.post(path, headers=headers, data=data)
+        self.assertEqual(201, response.status_code)
+
+        # Returned image entity should have a generated id and status
+        image = json.loads(response.text)
+        image_id = image['id']
+        self.assertEqual(image['status'], 'queued')
+        self.assertNotIn('size', image)
+
+        file_path = os.path.join(self.test_dir, 'fake_image')
+        with open(file_path, 'w') as fap:
+            fap.write('glance')
+
+        # Update locations for the queued image
+        path = self._url('/v2/images/%s' % image_id)
+        media_type = 'application/openstack-images-v2.1-json-patch'
+        headers = self._headers({'content-type': media_type})
+        data = json.dumps([{'op': 'replace', 'path': '/locations',
+                            'value': [{'url': 'file://' + file_path,
+                                       'metadata': {}}]}])
+        response = requests.patch(path, headers=headers, data=data)
+        self.assertEqual(200, response.status_code, response.text)
+
+        # The image size should be updated
+        path = self._url('/v2/images/%s' % image_id)
+        response = requests.get(path, headers=headers)
+        self.assertEqual(200, response.status_code)
+        image = json.loads(response.text)
+        self.assertEqual(image['size'], 6)
+
 
 class TestImageDirectURLVisibility(functional.FunctionalTest):
 

glance/tests/unit/test_store_image.py

@@ -44,6 +44,7 @@ class ImageStub(object):
         self.status = status
         self.locations = locations or []
         self.visibility = visibility
+        self.size = 1
 
     def delete(self):
         self.status = 'deleted'

glance/tests/unit/test_store_location.py

@@ -14,6 +14,7 @@
 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 #    License for the specific language governing permissions and limitations
 #    under the License.
+import mock
 
 from glance.common import exception
 from glance import context
@@ -427,3 +428,31 @@ class TestStoreLocation(base.StoreClearingUnitTest):
                               glance.store.get_store_from_scheme,
                               ctx,
                               store)
+
+    def test_add_location_for_image_without_size(self):
+        class FakeImageProxy():
+            size = None
+            context = None
+            store_api = mock.Mock()
+
+        def fake_get_size_from_backend(context, uri):
+            return 1
+
+        self.stubs.Set(glance.store, 'get_size_from_backend',
+                       fake_get_size_from_backend)
+        glance.store._check_image_location = mock.Mock()
+        loc1 = {'url': 'file:///fake1.img.tar.gz', 'metadata': {}}
+        loc2 = {'url': 'file:///fake2.img.tar.gz', 'metadata': {}}
+
+        # Test for insert location
+        image1 = FakeImageProxy()
+        locations = glance.store.StoreLocations(image1, [])
+        locations.insert(0, loc2)
+        self.assertEqual(image1.size, 1)
+
+        # Test for set_attr of _locations_proxy
+        image2 = FakeImageProxy()
+        locations = glance.store.StoreLocations(image2, [loc1])
+        locations[0] = loc2
+        self.assertTrue(loc2 in locations)
+        self.assertEqual(image2.size, 1)