单元测试模拟 urllib.request.urlretrieve() Python 3 和内部函数

Question

How can I mock or unit test a function/method that uses urllib.request.urlretrieve to save a file?

This is the part of code the I'm trying to to test:

from urllib import request
from config import cvs_site, proxy


class Collector(object):
    """Class Collector"""
    ...


def __init__(self, code_num=""):
    self.code_num = sec_id.upper()
    self.csv_file = "csv_files/code_num.csv"

    # load proxy if it is configured
    if proxy:
        proxies = {"http": proxy, "https": proxy, "ftp": proxy}
        proxy_connect = request.ProxyHandler(proxies)
        opener = request.build_opener(proxy_connect)
        request.install_opener(opener)

    def _collect_data():
        try:
            print("\nAccessing to retrieve CVS informations.")
            request.urlretrieve(cvs_site, self.cvs_file)
        except error.URLError as e:
            exit("\033[1;31m[ERROR]\033[1;00m {0}\n".format(e))

...

def some_function(self):
    _collect_data()
...

Should I test all internal functions (_functions())?

How to mock it?

Answer 1

To solve it I did some modifications on my code, to then created the test with mock.

REMARK: I'm still learning unit tests and mock, any new comments here is good because I'm not confident that I went to correct way :)

1. the function _collect_data() there is not necessities to be inside of __init__() , then I moved it to outside.

2. the _collect_data is a function with specific action, save the file, but need to return something to able to mock works with this.

3. the argument was moved from Class to function

The code new looks like it:

from config import proxy
from config import cvs_site
from urllib import request


class Collector(object):
    """Class Collector"""


def __init__(self):

    self.csv_file = "csv_files/code_num.csv"

    # load proxy if it is configured
    if proxy:
        proxies = {"http": proxy, "https": proxy, "ftp": proxy}
        proxy_connect = request.ProxyHandler(proxies)
        opener = request.build_opener(proxy_connect)
        request.install_opener(opener)


def _collect_data():
    try:
        print("\nAccessing to retrieve CVS informations.")
        return request.urlretrieve(cvs_site, self.cvs_file)
    except error.URLError as e:
        return "\033[1;31m[ERROR]\033[1;00m {0}\n".format(e)
...


def some_function(self, code_num=""):
    code_num = code_num.upper()
    self._collect_data()
...

To test this code I create this:

import unittest
import mock
from saassist.datacollector import Collector


class TestCollector(unittest.TestCase):
    def setUp(self):
        self.apar_test = Collector()

    @mock.patch("saassist.datacollector.request")
    def test_collect_data(self, mock_collect_data):

        mock_collect_data.urlretrieve.return_value = "File Collected OK"
        self.assertEqual("File Collected OK", self.apar_test._collect_data())

Well, I don't know which more could be tested for it, but to start I think it is good :)