def setUp(self):
self.apk = Apk(sha256='1234', pkg_name='apk1')
self.apk2 = Apk(sha256='5678', pkg_name='apk2')
self.dataset = Dataset(self.apk, self.apk2)
self.constructor_mock = mock.create_autospec(UrlConstructor)
self.constructor_mock.construct.side_effect = lambda apk: ' https://' + apk.sha256 # use sha256 of apk as download url, just for a test
self.dataset_downloader = DatasetDownloader(
self.dataset, url_constructor=self.constructor_mock, out_dir='out')
def run(input_file,
key,
number,
criteria,
metadata,
out_dir='azoo_dataset',
seed=None):
logging.info(f'START. APKS TO DOWNLOAD: {str(number)}')
source = None
try:
source = Source(input_file=input_file)
dataset = FilteredDataset(source, criteria)
random_subset = RandomPicker(seed=seed).get_random_subset(
dataset, number)
DatasetDownloader(dataset=random_subset, out_dir=out_dir,
key=key).download()
MetadataSaver(random_subset, out_dir).save(metadata)
except KeyboardInterrupt:
logging.info('Keyboard interrupt')
sys.exit()
finally:
if source:
source.input_file.close()
class DatasetDownloaderTest(unittest.TestCase):
def setUp(self):
self.apk = Apk(sha256='1234', pkg_name='apk1')
self.apk2 = Apk(sha256='5678', pkg_name='apk2')
self.dataset = Dataset(self.apk, self.apk2)
self.constructor_mock = mock.create_autospec(UrlConstructor)
self.constructor_mock.construct.side_effect = lambda apk: ' https://' + apk.sha256 # use sha256 of apk as download url, just for a test
self.dataset_downloader = DatasetDownloader(
self.dataset,
threads=4,
url_constructor=self.constructor_mock,
out_dir='out')
@mock.patch('requests.get')
@mock.patch("modules.services.dataset_downloader.open")
@mock.patch("modules.services.dataset_downloader.os")
def test_checks_if_out_dir_exists(self, os_mock, mock_open, mock_request):
self.dataset_downloader.download()
os_mock.path.exists.assert_any_call('out')
os_mock.path.exists.return_value = False
self.dataset_downloader.download()
os_mock.path.exists.assert_any_call('out')
os_mock.makedirs.assert_called_with('out')
@mock.patch("modules.services.dataset_downloader.os.makedirs")
@mock.patch("modules.services.dataset_downloader.open")
@mock.patch("modules.services.dataset_downloader.requests.get")
def test_downloads_apks(self, get_mock, open_mock, os_mock):
self.dataset_downloader.download()
self.assertTrue(
mock.call(self.apk) in
self.constructor_mock.construct.call_args_list)
self.assertTrue(
mock.call(self.apk2) in
self.constructor_mock.construct.call_args_list)
self.constructor_mock.reset_mock()
self.assertEqual(get_mock.call_count, 2)
self.assertTrue(
mock.call(r'out/apk1.apk', 'wb') in open_mock.call_args_list)
self.assertTrue(
mock.call(r'out/apk2.apk', 'wb') in open_mock.call_args_list)