test_datasets_utils.py 852 Bytes
Newer Older
1
import os
Francisco Massa's avatar
Francisco Massa committed
2
3
4
5
6
7
8
9
10
11
12
13
14
import shutil
import tempfile
import torch
import torchvision.datasets.utils as utils
import unittest


class Tester(unittest.TestCase):

    def test_download_url(self):
        temp_dir = tempfile.mkdtemp()
        url = "http://github.com/pytorch/vision/archive/master.zip"
        utils.download_url(url, temp_dir)
15
        assert not len(os.listdir(temp_dir)) == 0, 'The downloaded root directory is empty after download.'
Francisco Massa's avatar
Francisco Massa committed
16
17
18
19
20
21
        shutil.rmtree(temp_dir)

    def test_download_url_retry_http(self):
        temp_dir = tempfile.mkdtemp()
        url = "https://github.com/pytorch/vision/archive/master.zip"
        utils.download_url(url, temp_dir)
22
        assert not len(os.listdir(temp_dir)) == 0, 'The downloaded root directory is empty after download.'
Francisco Massa's avatar
Francisco Massa committed
23
24
25
26
27
        shutil.rmtree(temp_dir)


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