Fix multiple URLs handling - Merge pull request #30 from matslindh/cl…

…ean-up-url-handling

Fixes a couple of issues with multiple URLs, avoids always using the first URL sent in, and adds support for supplying a single string as the URL to use instead of requiring a list to be provided (since strings are lists as well, the old code will happily accept the string and then try to use each character as an URL).

The patch also picks a host for requests based on the first character of the image id, so we actually use most of the URLs supplied to the client. If there is a request without an image id, the URL defined first in the list is used.

imboclient/client.py

@@ -8,16 +8,19 @@
 from imboclient.header import authenticate
 from imboclient.url import image
 from imboclient.url import images
-from imboclient.url import user
+from imboclient.url import user as userUrl
 from imboclient.url import status
 from imboclient.url import metadata
 
+py_version = 3
+
 try:
     # py3
     from urllib.parse import urlparse
 except ImportError:
     # py2
     from urlparse import urlparse
+    py_version = 2
 
 class Client:
     def __init__(self, server_urls, public_key, private_key, version=None, user=None):
@@ -28,19 +31,19 @@ def __init__(self, server_urls, public_key, private_key, version=None, user=None
         self._user = user
 
     def metadata_url(self, image_identifier):
-        return metadata.UrlMetadata(self.server_urls[0], self._public_key, self._private_key, image_identifier, user=self._user)
+        return metadata.UrlMetadata(self._pick_url(image_identifier), self._public_key, self._private_key, image_identifier, user=self._user)
 
     def status_url(self):
-        return status.UrlStatus(self.server_urls[0], self._public_key, self._private_key)
+        return status.UrlStatus(self._pick_url(), self._public_key, self._private_key)
 
     def user_url(self):
-        return user.UrlUser(self.server_urls[0], self._public_key, self._private_key)
+        return userUrl.UrlUser(self._pick_url(), self._public_key, self._private_key)
 
     def images_url(self):
-        return images.UrlImages(self.server_urls[0], self._public_key, self._private_key, user=self._user)
+        return images.UrlImages(self._pick_url(), self._public_key, self._private_key, user=self._user)
 
     def image_url(self, image_identifier):
-        return image.UrlImage(self.server_urls[0], self._public_key, self._private_key, image_identifier, user=self._user)
+        return image.UrlImage(self._pick_url(image_identifier), self._public_key, self._private_key, image_identifier, user=self._user)
 
     def add_image(self, path):
         image_file_data = self._image_file_data(path)
@@ -145,7 +148,7 @@ def http_num_images(self):
         return 0
 
     def images(self, query = None):
-        images_url = images.UrlImages(self.server_urls[0], self._public_key, self._private_key)
+        images_url = images.UrlImages(self._pick_url(), self._public_key, self._private_key)
 
         if query:
             images_url.add_query(query)
@@ -204,19 +207,28 @@ def http_user_info(self):
 
         return self._wrap_result_json(http_user_info, [200], 'Failed getting user info.')
 
+    def _pick_url(self, key=None):
+        if key:
+            return self.server_urls[ord(key[0]) % len(self.server_urls)]
+        else:
+            return self.server_urls[0]
+
     def _image_file_data(self, path):
         return open(path, 'rb').read()
 
     def _parse_urls(self, urls):
         def should_remove_port(self, url_parts):
-            return parts.port and (parts.scheme == 'http' and parts.port == 80 or (parts.scheme == 'https' and parts.port == 443))
+            return url_parts.port and (url_parts.scheme == 'http' and url_parts.port == 80 or (url_parts.scheme == 'https' and url_parts.port == 443))
 
         parsed_urls = []
         pattern = re.compile("https?://")
 
+        if isinstance(urls, str) or (py_version == 2 and isinstance(urls, basestring)):
+            urls = [urls]
+
         for url in urls:
             if not pattern.match(url):
-                parsed_urls.append('http://' + url)
+                url = 'http://' + url
 
             parts = urlparse(url)
 
@@ -239,10 +251,6 @@ def _validate_local_file(self, path):
     def _generate_image_identifier(self, content):
         return hashlib.md5(content).hexdigest()
 
-    def _host_for_image_identifier(self, image_identifier):
-        dec = int(image_identifier[0] + image_identifier[1], 16)
-        return self.server_urls[dec % len(self.server_urls)]
-
     def _wrap_result_json(self, function, success_status_codes, error):
         try:
             response = self._wrap_result(function, success_status_codes, error)

imboclient/test/unit/test_client.py

@@ -51,6 +51,23 @@ def test_server_urls_port_explicit_without_protocol(self):
         self._client = imbo.Client(['imbo.local:8000'], 'public', 'private')
         assert self._client.server_urls[0] == 'http://imbo.local:8000'
 
+    def test_server_urls_as_string(self):
+        client = imbo.Client('imbo.local', 'public', 'private')
+        assert client.server_urls[0] == 'http://imbo.local'
+
+    def test_server_urls_from_identifiers(self):
+        hosts = ('imbo.local', 'imbo.local2', )
+        client = imbo.Client(hosts, 'public', 'private')
+
+        host1 = str(client.image_url('foo')).split('/')[2]
+
+        # the default method uses the ord() value of the chars, so increase by one to get a different host
+        host2 = str(client.image_url('goo')).split('/')[2]
+
+        assert host1 != host2
+        assert host1 in hosts
+        assert host2 in hosts
+
     @patch('imboclient.url.status.UrlStatus')
     def test_status_url(self, mocked_url_status):
         status_url = self._client.status_url()