Merge pull request #1989 from marxide/clear-session-header-range

Delete range from session request header after use

Author: bsipocz

astroquery/query.py

@@ -210,7 +210,7 @@ def _request(self, method, url,
                  files=None, save=False, savedir='', timeout=None, cache=True,
                  stream=False, auth=None, continuation=True, verify=True,
                  allow_redirects=True,
-                 json=None):
+                 json=None, return_response_on_save=False):
         """
         A generic HTTP request method, similar to `requests.Session.request`
         but with added caching-related tools
@@ -249,14 +249,21 @@ def _request(self, method, url,
             parameter will try to continue the download where it left off.
             See `_download_file`.
         stream : bool
+        return_response_on_save : bool
+            If ``save``, also return the server response. The default is to only
+            return the local file path.
 
         Returns
         -------
         response : `requests.Response`
             The response from the server if ``save`` is False
         local_filepath : list
             a list of strings containing the downloaded local paths if ``save``
-            is True
+            is True and ``return_response_on_save`` is False.
+        (local_filepath, response) : tuple(list, `requests.Response`)
+            a tuple containing a list of strings containing the downloaded local paths,
+            and the server response object, if ``save`` is True and ``return_response_on_save``
+            is True.
         """
         req_kwargs = dict(
             params=params,
@@ -274,11 +281,14 @@ def _request(self, method, url,
                 local_filename = local_filename.replace(':', '_')
             local_filepath = os.path.join(savedir or self.cache_location or '.', local_filename)
 
-            self._download_file(url, local_filepath, cache=cache,
-                                continuation=continuation, method=method,
-                                allow_redirects=allow_redirects,
-                                auth=auth, **req_kwargs)
-            return local_filepath
+            response = self._download_file(url, local_filepath, cache=cache,
+                                           continuation=continuation, method=method,
+                                           allow_redirects=allow_redirects,
+                                           auth=auth, **req_kwargs)
+            if return_response_on_save:
+                return local_filepath, response
+            else:
+                return local_filepath
         else:
             query = AstroQuery(method, url, **req_kwargs)
             if ((self.cache_location is None) or (not self._cache_active) or (not cache)):
@@ -369,6 +379,7 @@ def _download_file(self, url, local_filepath, timeout=None, auth=None,
                                                  timeout=timeout, stream=True,
                                                  auth=auth, **kwargs)
                 response.raise_for_status()
+                del self._session.headers['Range']
 
         elif cache and os.path.exists(local_filepath):
             if length is not None:

astroquery/tests/test_resume.py

@@ -9,18 +9,17 @@
 
 
 @pytest.mark.skipif('not ACTIVE_HTTPBIN')
-def test_resume():
+def test_resume(length=2048, partial_length=1024, qu=None):
     # Test that a resumed query will finish
 
-    length = 2048
-
     # 'range' will return abcd...xyz repeating
     target_url = 'http://127.0.0.1:5000/range/{0}'.format(length)
 
     # simple check: make sure the server does what's expected
     assert len(requests.get(target_url).content) == length
 
-    qu = query.BaseQuery()
+    if qu is None:
+        qu = query.BaseQuery()
 
     result_1 = qu._request('GET', target_url, save=True)
     # now the full file is written, so we have to delete parts of it
@@ -29,17 +28,29 @@ def test_resume():
         data = fh.read()
     with open(result_1, 'wb') as fh:
         # overwrite with a partial file
-        fh.write(data[:1024])
+        fh.write(data[:partial_length])
     with open(result_1, 'rb') as fh:
         data = fh.read()
-        assert len(data) == 1024
+        assert len(data) == partial_length
 
-    result_2 = qu._request('GET', target_url, save=True, continuation=True)
+    result_2, response = qu._request('GET', target_url, save=True, continuation=True,
+                                     return_response_on_save=True)
 
-    assert 'range' in qu._session.headers
-    assert qu._session.headers['range'] == 'bytes={0}-{1}'.format(1024, length-1)
+    assert 'content-range' in response.headers
+    assert response.headers['content-range'] == 'bytes {0}-{1}/{2}'.format(
+        partial_length, length-1, length
+    )
 
     with open(result_2, 'rb') as fh:
         data = fh.read()
     assert len(data) == length
-    assert data == (string.ascii_lowercase*80)[:length].encode('ascii')
+    assert data == (string.ascii_lowercase*(length//26+1))[:length].encode('ascii')
+
+
+@pytest.mark.skipif('not ACTIVE_HTTPBIN')
+def test_resume_consecutive():
+    # Test that consecutive resumed queries request the correct content range and finish
+    qu = query.BaseQuery()
+
+    test_resume(length=2048, partial_length=1024, qu=qu)
+    test_resume(length=2048, partial_length=512, qu=qu)