From 18f38ff8e83305ac0d0738819d596ba0e5e4be2d Mon Sep 17 00:00:00 2001 From: Marek Goldmann Date: Thu, 6 Apr 2017 14:42:05 +0200 Subject: [PATCH] Better handling of fetching files and cacher url --- dogen/cli.py | 2 +- dogen/generator.py | 14 +++++++++++--- tests/test_unit_generate_configuration.py | 12 +++++------- tests/test_unit_generate_handle_files.py | 16 ++++++++++++---- 4 files changed, 29 insertions(+), 15 deletions(-) diff --git a/dogen/cli.py b/dogen/cli.py index 01680df..2e8930d 100644 --- a/dogen/cli.py +++ b/dogen/cli.py @@ -25,7 +25,7 @@ def error(self, message): class CLI(object): def __init__(self): - self.log = logging.getLogger() + self.log = logging.getLogger("dogen") handler = logging.StreamHandler() formatter = logging.Formatter( '%(asctime)s %(name)-12s %(levelname)-8s %(message)s') diff --git a/dogen/generator.py b/dogen/generator.py index 78218fb..6d0793b 100644 --- a/dogen/generator.py +++ b/dogen/generator.py @@ -54,8 +54,14 @@ def _fetch_file(self, location, output=None): self.log.info("Fetched file will be saved as '%s'..." % os.path.basename(output)) + r = requests.get(location, verify=self.ssl_verify, stream=True) + + if r.status_code != 200: + raise Exception("Could not download file from %s" % location) + with open(output, 'wb') as f: - f.write(requests.get(location, verify=self.ssl_verify).content) + for chunk in r.iter_content(chunk_size=1024): + f.write(chunk) return output @@ -294,11 +300,13 @@ def handle_sources(self): if not passed: sources_cache = os.environ.get("DOGEN_SOURCES_CACHE") + if sources_cache: - self.log.info("Using '%s' as cached location for sources" % sources_cache) - url = "%s/%s" % (sources_cache, basename) + url = sources_cache.replace('#hash#', source['md5sum']).replace('#algorithm#', 'md5') + self.log.info("Using '%s' as cached location for artifact" % url) self._fetch_file(url, filename) + self.check_sum(filename, source['md5sum']) self.cfg['artifacts'][target] = "md5:%s" % source['md5sum'] diff --git a/tests/test_unit_generate_configuration.py b/tests/test_unit_generate_configuration.py index 1799ac1..9c01463 100644 --- a/tests/test_unit_generate_configuration.py +++ b/tests/test_unit_generate_configuration.py @@ -1,8 +1,6 @@ import argparse import unittest import mock -import six -import tarfile import os import tempfile @@ -35,7 +33,7 @@ def test_default_values(self): self.assertEqual(self.generator.scripts_path, None) self.assertEqual(self.generator.additional_scripts, None) self.assertEqual(self.generator.without_sources, False) - # Set to True in the configure() method later + # Set to True in the configure() method later self.assertEqual(self.generator.ssl_verify, None) def test_fail_if_version_mismatch(self): @@ -319,14 +317,14 @@ def test_handling_sources_with_cache_url(self, mock_fetch_file, mock_check_sum): with self.descriptor as f: f.write("sources:\n - url: http://somehost.com/file.zip\n md5sum: e9013fc202c87be48e3b302df10efc4b".encode()) - k = mock.patch.dict(os.environ, {'DOGEN_SOURCES_CACHE':'http://cache/context'}) + k = mock.patch.dict(os.environ, {'DOGEN_SOURCES_CACHE':'http://cache/get?#algorithm#=#hash#'}) k.start() generator = Generator(self.log, self.args) generator.configure() generator.handle_sources() k.stop() - mock_fetch_file.assert_called_with('http://cache/context/file.zip', 'target/file.zip') + mock_fetch_file.assert_called_with('http://cache/get?md5=e9013fc202c87be48e3b302df10efc4b', 'target/file.zip') @mock.patch('dogen.generator.Generator.check_sum') @mock.patch('dogen.generator.Generator._fetch_file') @@ -334,11 +332,11 @@ def test_handling_sources_with_cache_url_and_target_filename(self, mock_fetch_fi with self.descriptor as f: f.write("sources:\n - url: http://somehost.com/file.zip\n md5sum: e9013fc202c87be48e3b302df10efc4b\n target: target.zip".encode()) - k = mock.patch.dict(os.environ, {'DOGEN_SOURCES_CACHE':'http://cache/context'}) + k = mock.patch.dict(os.environ, {'DOGEN_SOURCES_CACHE':'http://cache/get?#algorithm#=#hash#'}) k.start() generator = Generator(self.log, self.args) generator.configure() generator.handle_sources() k.stop() - mock_fetch_file.assert_called_with('http://cache/context/file.zip', 'target/target.zip') + mock_fetch_file.assert_called_with('http://cache/get?md5=e9013fc202c87be48e3b302df10efc4b', 'target/target.zip') diff --git a/tests/test_unit_generate_handle_files.py b/tests/test_unit_generate_handle_files.py index 749eee1..eb85ab9 100644 --- a/tests/test_unit_generate_handle_files.py +++ b/tests/test_unit_generate_handle_files.py @@ -34,11 +34,15 @@ def setUp(self): @mock.patch('dogen.generator.requests.get') def test_fetching_with_filename(self, mock_requests): - mock_requests.return_value.content = "file-content" + def iter_content(**args): + return ["file-content"] + + mock_requests.return_value.status_code = 200 + mock_requests.return_value.iter_content = iter_content with mock.patch.object(six.moves.builtins, 'open', mock.mock_open()) as mock_file: self.assertEqual(self.generator._fetch_file("https://host/file.tmp", "some-file"), "some-file") - mock_requests.assert_called_with('https://host/file.tmp', verify=None) + mock_requests.assert_called_with('https://host/file.tmp', verify=None, stream=True) mock_file().write.assert_called_once_with("file-content") self.log.info.assert_any_call("Fetching 'https://host/file.tmp' file...") @@ -48,12 +52,16 @@ def test_fetching_with_filename(self, mock_requests): @mock.patch('dogen.generator.tempfile.mktemp', return_value="tmpfile") @mock.patch('dogen.generator.requests.get') def test_fetching_with_tmpfile(self, mock_requests, mock_tempfile): - mock_requests.return_value.content = "file-content" + def iter_content(**args): + return ["file-content"] + + mock_requests.return_value.status_code = 200 + mock_requests.return_value.iter_content = iter_content with mock.patch.object(six.moves.builtins, 'open', mock.mock_open()) as mock_file: self.assertEqual(self.generator._fetch_file("https://host/file.tmp"), "tmpfile") mock_tempfile.assert_called_with("-dogen") - mock_requests.assert_called_with('https://host/file.tmp', verify=None) + mock_requests.assert_called_with('https://host/file.tmp', verify=None, stream=True) mock_file().write.assert_called_once_with("file-content") self.log.info.assert_any_call("Fetching 'https://host/file.tmp' file...")