diff --git a/bonobo/commands/download.py b/bonobo/commands/download.py index 0106645..fd51951 100644 --- a/bonobo/commands/download.py +++ b/bonobo/commands/download.py @@ -29,8 +29,7 @@ def execute(path, *args, **kwargs): raise ValueError('download command currently supports examples only') examples_path = re.sub('^examples/', '', path) output_path = bonobo.get_examples_path(examples_path) - response = _open_url(EXAMPLES_BASE_URL + examples_path) - with open(output_path, 'wb') as fout: + with _open_url(EXAMPLES_BASE_URL + examples_path) as response, open(output_path, 'wb') as fout: _write_response(response, fout) print('saved to {}'.format(output_path)) diff --git a/tests/test_commands.py b/tests/test_commands.py index f68657c..c78fa5f 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -164,6 +164,12 @@ def test_download_works_for_examples(runner): def iter_content(self, *args, **kwargs): return [expected_bytes] + def __enter__(self): + return self + + def __exit__(self, *args, **kwargs): + pass + fout = io.BytesIO() fout.close = lambda: None