Rewind original stream body when refreshing.

When refreshing credentials, the original request is re-sent
after the credentials are refreshed.

If the body of that request is a stream, the stream contents
are read in the initial request, and the stream must be rewound
before the request is re-sent. Otherwise, the original message
body will be different (because stream data was skipped).

1 files changed, 38 insertions, 5 deletions

diff --git a/tests/test_file.py b/tests/test_file.py
index 8efe6cf..89c3e02 100644
--- a/tests/test_file.py
+++ b/tests/test_file.py
@@ -32,12 +32,15 @@ import tempfile
 import unittest
 
 from .http_mock import HttpMockSequence
+import six
+
 from oauth2client import file
 from oauth2client import locked_file
 from oauth2client import multistore_file
 from oauth2client import util
 from oauth2client.client import AccessTokenCredentials
 from oauth2client.client import OAuth2Credentials
+from six.moves import http_client
 try:
   # Python2
   from future_builtins import oct
@@ -154,15 +157,17 @@ class OAuth2ClientFileTests(unittest.TestCase):
     access_token = '1/3w'
     token_response = {'access_token': access_token, 'expires_in': 3600}
     http = HttpMockSequence([
-        ({'status': '401'}, b'Initial token expired'),
-        ({'status': '401'}, b'Store token expired'),
-        ({'status': '200'}, json.dumps(token_response).encode('utf-8')),
-        ({'status': '200'}, b'Valid response to original request')
+        ({'status': str(http_client.UNAUTHORIZED)}, b'Initial token expired'),
+        ({'status': str(http_client.UNAUTHORIZED)}, b'Store token expired'),
+        ({'status': str(http_client.OK)},
+         json.dumps(token_response).encode('utf-8')),
+        ({'status': str(http_client.OK)},
+         b'Valid response to original request')
     ])
 
     credentials.authorize(http)
     http.request('https://example.com')
-    self.assertEquals(credentials.access_token, access_token)
+    self.assertEqual(credentials.access_token, access_token)
 
   def test_token_refresh_good_store(self):
     expiration = datetime.datetime.utcnow() + datetime.timedelta(minutes=15)
@@ -178,6 +183,34 @@ class OAuth2ClientFileTests(unittest.TestCase):
     credentials._refresh(lambda x: x)
     self.assertEquals(credentials.access_token, 'bar')
 
+  def test_token_refresh_stream_body(self):
+    expiration = datetime.datetime.utcnow() + datetime.timedelta(minutes=15)
+    credentials = self.create_test_credentials(expiration=expiration)
+
+    s = file.Storage(FILENAME)
+    s.put(credentials)
+    credentials = s.get()
+    new_cred = copy.copy(credentials)
+    new_cred.access_token = 'bar'
+    s.put(new_cred)
+
+    valid_access_token = '1/3w'
+    token_response = {'access_token': valid_access_token, 'expires_in': 3600}
+    http = HttpMockSequence([
+        ({'status': str(http_client.UNAUTHORIZED)}, b'Initial token expired'),
+        ({'status': str(http_client.UNAUTHORIZED)}, b'Store token expired'),
+        ({'status': str(http_client.OK)},
+         json.dumps(token_response).encode('utf-8')),
+        ({'status': str(http_client.OK)}, 'echo_request_body')
+    ])
+
+    body = six.StringIO('streaming body')
+
+    credentials.authorize(http)
+    _, content = http.request('https://example.com', body=body)
+    self.assertEqual(content, 'streaming body')
+    self.assertEqual(credentials.access_token, valid_access_token)
+
   def test_credentials_delete(self):
     credentials = self.create_test_credentials()