Merge "Wrap raw iterators to ensure we send entire contents to server"

3 files changed, 34 insertions, 1 deletions

diff --git a/swiftclient/client.py b/swiftclient/client.py
index 34edd11..aba986e 100644
--- a/swiftclient/client.py
+++ b/swiftclient/client.py
@@ -32,7 +32,7 @@ import six
 from swiftclient import version as swiftclient_version
 from swiftclient.exceptions import ClientException
 from swiftclient.utils import (
-    LengthWrapper, ReadableToIterable, parse_api_response)
+    iter_wrapper, LengthWrapper, ReadableToIterable, parse_api_response)
 
 # Default is 100, increase to 256
 http_client._MAXHEADERS = 256
@@ -1126,6 +1126,10 @@ def put_object(url, token=None, container=None, name=None, contents=None,
             warn_msg = ('%s object has no "read" method, ignoring chunk_size'
                         % type(contents).__name__)
             warnings.warn(warn_msg, stacklevel=2)
+        # Match requests's is_stream test
+        if hasattr(contents, '__iter__') and not isinstance(contents, (
+                six.text_type, six.binary_type, list, tuple, dict)):
+            contents = iter_wrapper(contents)
         conn.request('PUT', path, contents, headers)
 
     resp = conn.getresponse()
diff --git a/swiftclient/utils.py b/swiftclient/utils.py
index 742ec06..9d94b6b 100644
--- a/swiftclient/utils.py
+++ b/swiftclient/utils.py
@@ -231,3 +231,12 @@ class LengthWrapper(object):
             self.md5sum.update(chunk.encode())
 
         return chunk
+
+
+def iter_wrapper(iterable):
+    for chunk in iterable:
+        if len(chunk) == 0:
+            # If we emit an empty chunk, requests will go ahead and send it,
+            # causing the server to close the connection
+            continue
+        yield chunk
diff --git a/tests/unit/test_swiftclient.py b/tests/unit/test_swiftclient.py
index e837288..60f65c9 100644
--- a/tests/unit/test_swiftclient.py
+++ b/tests/unit/test_swiftclient.py
@@ -984,6 +984,26 @@ class TestPutObject(MockHttpTest):
                 data += chunk
             self.assertEqual(data, raw_data)
 
+    def test_iter_upload(self):
+        def data():
+            for chunk in ('foo', '', 'bar'):
+                yield chunk
+        conn = c.http_connection(u'http://www.test.com/')
+        resp = MockHttpResponse(status=200)
+        conn[1].getresponse = resp.fake_response
+        conn[1]._request = resp._fake_request
+
+        c.put_object(url='http://www.test.com', http_conn=conn,
+                     contents=data())
+        req_headers = resp.requests_params['headers']
+        self.assertNotIn('Content-Length', req_headers)
+        req_data = resp.requests_params['data']
+        self.assertTrue(hasattr(req_data, '__iter__'))
+        # If we emit an empty chunk, requests will go ahead and send it,
+        # causing the server to close the connection. So make sure we don't
+        # do that.
+        self.assertEqual(['foo', 'bar'], list(req_data))
+
     def test_md5_mismatch(self):
         conn = c.http_connection('http://www.test.com')
         resp = MockHttpResponse(status=200, verify=True,