diff --git a/config.yaml b/config.yaml index a2d20a0..81e0e30 100755 --- a/config.yaml +++ b/config.yaml @@ -2,6 +2,7 @@ service: logging_level: $LOGGING_LEVEL_ROOT|DEBUG # Logging level for service logger response_file_extension: $RESPONSE_FILE_EXTENSION|"json.gz" # Extension for files uploaded to storage response_folder: $RESPONSE_FOLDER|null + download_strategy: single probing_webserver: host: $PROBING_WEBSERVER_HOST|"0.0.0.0" # Probe webserver address diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index 25bc1e5..8c2d7a2 100644 --- a/pyinfra/visitor.py +++ b/pyinfra/visitor.py @@ -300,7 +300,7 @@ def get_download_strategy(download_strategy_type=None): "single": SingleDownloadStrategy(), "multi": MultiDownloadStrategy(), } - return download_strategies.get(download_strategy_type or CONFIG.download_strategy, SingleDownloadStrategy()) + return download_strategies.get(download_strategy_type or CONFIG.service.download_strategy, SingleDownloadStrategy()) class DownloadStrategy(abc.ABC): diff --git a/test/unit_tests/queue_visitor_test.py b/test/unit_tests/queue_visitor_test.py index 87a184d..ce536aa 100644 --- a/test/unit_tests/queue_visitor_test.py +++ b/test/unit_tests/queue_visitor_test.py @@ -20,8 +20,8 @@ class TestVisitor: ): storage.clear_bucket(bucket_name) storage.put_object(**get_object_descriptor(body), data=pack_for_upload(b"content")) - data_received = visitor.load_data(body) - assert {"data": b"content", "metadata": {}} == data_received + data_received = list(visitor.load_data(body)) + assert [{"data": b"content", "metadata": {}}] == data_received @pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session") def test_visitor_pulls_and_processes_data(self, visitor, body, storage, bucket_name):