CODEWITHSUNDEEP
pythonpython-requestsmocking
Kim Stacks
Kim Stacks

Reputation: 10812

How to write mock for file download for unit tests

I am using requests_mock as the library to test code I write that consumes endpoints.

This is my download_file function

def download_as_file(url: str, auth=(), file_path="", attempts=2):
    """Downloads a URL content into a file (with large file support by streaming)

    :param url: URL to download
    :param auth: tuple containing credentials to access the url
    :param file_path: Local file name to contain the data downloaded
    :param attempts: Number of attempts
    :return: New file path. Empty string if the download failed
    """
    if not file_path:
        file_path = os.path.realpath(os.path.basename(url))
    logger.info(f"Downloading {url} content to {file_path}")
    url_sections = urlparse(url)
    if not url_sections.scheme:
        logger.debug("The given url is missing a scheme. Adding http scheme")
        url = f"https://{url}"
        logger.debug(f"New url: {url}")
    for attempt in range(1, attempts + 1):
        try:
            if attempt > 1:
                time.sleep(10)  # 10 seconds wait time between downloads
            with requests.get(url, auth=auth, stream=True) as response:
                response.raise_for_status()
                with open(file_path, "wb") as out_file:
                    for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB chunks
                        out_file.write(chunk)
                logger.info("Download finished successfully")
                return (response, file_path)
        except Exception as ex:
            logger.error(f"Attempt #{attempt} failed with error: {ex}")
    return None

How do I test this with requests_mock?

I know the documentation for requests_mock covers the following at https://requests-mock.readthedocs.io/en/latest/response.html#registering-responses

To specify the body of the response there are a number of options that depend on the format that you wish to return.

  • json: A python object that will be converted to a JSON string.
  • text: A unicode string. This is typically what you will want to use for regular textual content.
  • content: A byte string. This should be used for including binary data in responses.
  • body: A file like object that contains a .read() function.
  • raw: A prepopulated urllib3.response.HTTPResponse to be returned.
  • exc: An exception that will be raised instead of returning a response.

Is it body or content? And how do I write it?

Most endpoints give me json, but I have this particular use case of downloading a .xlsx file, so I want to write a test case for consuming it.

Upvotes: 1

Views: 2826

Answers (1)

Kim Stacks
Kim Stacks

Reputation: 10812

Short answer: Choose body

Long answer:

with requests_mock.mock() as m:
    current_folder = os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(
        current_folder, filename
    )
    with open(path, "rb") as the_file:
        m.get(url_to_mock, body=the_file)

Upvotes: 3

Related Questions