CODEWITHSUNDEEP
pythonpython-2.7unit-testingdownloadmocking
Ari Linn
Ari Linn

Reputation: 207

Is there any way to mock file download in Python?

I'm using Python 2.7 with Mock. I have a method that takes a url and downloads it into a temporary file, then renames file according to business logic rules. I want to test this renaming logic, but first I must mock the file download part. And here is where I'm stuck. I'm using urllib2.urlopen and its read(chunkSize) method in infinite cycle, checking if read(chunkSize) returns some value. While this approach works in real life, and the response is eventually read up to the end where read(chunkSize) does not return anything, I get an infinite cycle when mocking. read(chunkSize) seems to always have a result. How do I get the cycle to stop once the contents of the response has been read? Here's my code:

import urllib2
from contextlib import closing
import unittest
import mock

def Method(url, temppath):
    chunkSize = 16 * 1024
    request = urllib2.Request(url)
    with closing(urllib2.urlopen(request, timeout = 5)) as response:
        with open(temppath, 'wb') as largeFile:
            while True:
                chunk = response.read(chunkSize)
                # print chunk # <- this will endlessly produce '0123456' when tested by test_Method in MyTestCase
                if not chunk:
                    break
                largeFile.write(chunk)
    # rename file from temppath to something new

class MyTestCase(unittest.TestCase):
    @mock.patch('urllib2.urlopen', autospec=True)
    @mock.patch('__main__.open', create=True)
    def test_Method(self, mock_open, mock_urlopen):
        mock_urlopen.return_value.read.return_value = b'0123456'
        Method('http://a.bcd/img.png', 'a:\\b\\1234567890.tmp')

if __name__ == '__main__':
    unittest.main()

Upvotes: 2

Views: 4269

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1122532

Assign a list of values to return to the side_effect attribute:

mock_urlopen.return_value.read.side_effect = [b'0123456', b'']

The mock will iterate over the values for each read() call, so the last call returning an empty bytes object terminates your loop.

Upvotes: 4

Related Questions