python unittest howto

I`d like to know how I could unit-test the following module.

def download_distribution(url, tempdir):
    """ Method which downloads the distribution from PyPI """
    print "Attempting to download from %s" % (url,)

    try:
        url_handler = urllib2.urlopen(url)
        distribution_contents = url_handler.read()
        url_handler.close()

        filename = get_file_name(url)

        file_handler = open(os.path.join(tempdir, filename), "w")
        file_handler.write(distribution_contents)
 开发者_如何学JAVA       file_handler.close()
        return True

    except ValueError, IOError:
        return False

Unit test propositioners will tell you that unit tests should be self contained, that is, they should not access the network or the filesystem (especially not in writing mode). Network and filesystem tests are beyond the scope of unit tests (though you might subject them to integration tests).

Speaking generally, for such a case, I'd extract the urllib and file-writing codes to separate functions (which would not be unit-tested), and inject mock-functions during unit testing.

I.e. (slightly abbreviated for better reading):

def get_web_content(url):
    # Extracted code
    url_handler = urllib2.urlopen(url)
    content = url_handler.read()
    url_handler.close()
    return content

def write_to_file(content, filename, tmpdir):
    # Extracted code
    file_handler = open(os.path.join(tempdir, filename), "w")
    file_handler.write(content)
    file_handler.close()

def download_distribution(url, tempdir):
    # Original code, after extractions
    distribution_contents = get_web_content(url)
    filename = get_file_name(url)
    write_to_file(distribution_contents, filename, tmpdir)
    return True

And, on the test file:

import module_I_want_to_test

def mock_web_content(url):
    return """Some fake content, useful for testing"""
def mock_write_to_file(content, filename, tmpdir):
    # In this case, do nothing, as we don't do filesystem meddling while unit testing
    pass

module_I_want_to_test.get_web_content = mock_web_content
module_I_want_to_test.write_to_file = mock_write_to_file

class SomeTests(unittest.Testcase):
    # And so on...

And then I second Daniel's suggestion, you should read some more in-depth material on unit testing.

Vague question. If you're just looking for a primer for unit testing in general with a Python slant, I recommend Mark Pilgrim's "Dive Into Python" which has a chapter on unit testing with Python. Otherwise you need to clear up what specific issues you are having testing that code.

To mock urllopen you can pre fetch some examples that you can then use in your unittests. Here's an example to get you started:

def urlopen(url):
    urlclean = url[:url.find('?')] # ignore GET parameters
    files = {
        'http://example.com/foo.xml': 'foo.xml',
        'http://example.com/bar.xml': 'bar.xml',
    }
    return file(files[urlclean])
yourmodule.urllib.urlopen = urlopen