Fix issue 10817 - Fix urlretrieve function to raise ContentTooShortError

orsenthil · orsenthil · commit ce260142c6c6 · 2011-11-01T01:35:17.000+08:00
even when reporthook is None. Patch by Jyrki Pulliainen.
diff --git a/Lib/test/test_urllib.py b/Lib/test/test_urllib.py
@@ -36,6 +36,44 @@ def urlopen(url, data=None, proxies=None):
     else:
         return opener.open(url, data)
 
+
+class FakeHTTPMixin(object):
+    def fakehttp(self, fakedata):
+        class FakeSocket(io.BytesIO):
+            io_refs = 1
+
+            def sendall(self, str):
+                pass
+
+            def makefile(self, *args, **kwds):
+                self.io_refs += 1
+                return self
+
+            def read(self, amt=None):
+                if self.closed:
+                    return b""
+                return io.BytesIO.read(self, amt)
+
+            def readline(self, length=None):
+                if self.closed:
+                    return b""
+                return io.BytesIO.readline(self, length)
+
+            def close(self):
+                self.io_refs -= 1
+                if self.io_refs == 0:
+                    io.BytesIO.close(self)
+
+        class FakeHTTPConnection(http.client.HTTPConnection):
+            def connect(self):
+                self.sock = FakeSocket(fakedata)
+        self._connection_class = http.client.HTTPConnection
+        http.client.HTTPConnection = FakeHTTPConnection
+
+    def unfakehttp(self):
+        http.client.HTTPConnection = self._connection_class
+
+
 class urlopen_FileTests(unittest.TestCase):
     """Test urlopen() opening a temporary file.
 
@@ -139,35 +177,9 @@ def test_getproxies_environment_keep_no_proxies(self):
         self.env.set('NO_PROXY', 'localhost, anotherdomain.com, newdomain.com')
         self.assertTrue(urllib.request.proxy_bypass_environment('anotherdomain.com'))
 
-class urlopen_HttpTests(unittest.TestCase):
+class urlopen_HttpTests(unittest.TestCase, FakeHTTPMixin):
     """Test urlopen() opening a fake http connection."""
 
-    def fakehttp(self, fakedata):
-        class FakeSocket(io.BytesIO):
-            io_refs = 1
-            def sendall(self, str): pass
-            def makefile(self, *args, **kwds):
-                self.io_refs += 1
-                return self
-            def read(self, amt=None):
-                if self.closed: return b""
-                return io.BytesIO.read(self, amt)
-            def readline(self, length=None):
-                if self.closed: return b""
-                return io.BytesIO.readline(self, length)
-            def close(self):
-                self.io_refs -= 1
-                if self.io_refs == 0:
-                    io.BytesIO.close(self)
-        class FakeHTTPConnection(http.client.HTTPConnection):
-            def connect(self):
-                self.sock = FakeSocket(fakedata)
-        self._connection_class = http.client.HTTPConnection
-        http.client.HTTPConnection = FakeHTTPConnection
-
-    def unfakehttp(self):
-        http.client.HTTPConnection = self._connection_class
-
     def check_read(self, ver):
         self.fakehttp(b"HTTP/" + ver + b" 200 OK\r\n\r\nHello!")
         try:
@@ -394,6 +406,48 @@ def hooktester(count, block_size, total_size, _report=report):
         self.assertEqual(report[0][1], 8192)
         self.assertEqual(report[0][2], 8193)
 
+
+class urlretrieve_HttpTests(unittest.TestCase, FakeHTTPMixin):
+    """Test urllib.urlretrieve() using fake http connections"""
+
+    def test_short_content_raises_ContentTooShortError(self):
+        self.fakehttp(b'''HTTP/1.1 200 OK
+Date: Wed, 02 Jan 2008 03:03:54 GMT
+Server: Apache/1.3.33 (Debian GNU/Linux) mod_ssl/2.8.22 OpenSSL/0.9.7e
+Connection: close
+Content-Length: 100
+Content-Type: text/html; charset=iso-8859-1
+
+FF
+''')
+
+        def _reporthook(par1, par2, par3):
+            pass
+
+        with self.assertRaises(urllib.error.ContentTooShortError):
+            try:
+                urllib.request.urlretrieve('http://example.com/',
+                                           reporthook=_reporthook)
+            finally:
+                self.unfakehttp()
+
+    def test_short_content_raises_ContentTooShortError_without_reporthook(self):
+        self.fakehttp(b'''HTTP/1.1 200 OK
+Date: Wed, 02 Jan 2008 03:03:54 GMT
+Server: Apache/1.3.33 (Debian GNU/Linux) mod_ssl/2.8.22 OpenSSL/0.9.7e
+Connection: close
+Content-Length: 100
+Content-Type: text/html; charset=iso-8859-1
+
+FF
+''')
+        with self.assertRaises(urllib.error.ContentTooShortError):
+            try:
+                urllib.request.urlretrieve('http://example.com/')
+            finally:
+                self.unfakehttp()
+
+
 class QuotingTests(unittest.TestCase):
     """Tests for urllib.quote() and urllib.quote_plus()
 
@@ -1164,6 +1218,7 @@ def test_main():
         urlopen_FileTests,
         urlopen_HttpTests,
         urlretrieve_FileTests,
+        urlretrieve_HttpTests,
         ProxyTests,
         QuotingTests,
         UnquotingTests,
diff --git a/Lib/urllib/request.py b/Lib/urllib/request.py
@@ -1589,9 +1589,9 @@ def retrieve(self, url, filename=None, reporthook=None, data=None):
                 size = -1
                 read = 0
                 blocknum = 0
+                if "content-length" in headers:
+                    size = int(headers["Content-Length"])
                 if reporthook:
-                    if "content-length" in headers:
-                        size = int(headers["Content-Length"])
                     reporthook(blocknum, bs, size)
                 while 1:
                     block = fp.read(bs)
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -66,6 +66,9 @@ Core and Builtins
 Library
 -------
 
+- Issue 10817: Fix urlretrieve function to raise ContentTooShortError even
+  when reporthook is None. Patch by Jyrki Pulliainen.
+
 - Issue 13296: Fix IDLE to clear compile __future__ flags on shell restart.
   (Patch by Roger Serwy)
 
