From 9130b9fc753c63bfe4abd3cc26b31caf5c5bf2cb Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Mon, 14 Mar 2022 19:05:59 +0100 Subject: [PATCH] test refac --- pyinfra/test/unit_tests/conftest.py | 4 ++-- pyinfra/test/unit_tests/queue_visitor_test.py | 10 ++++------ pyinfra/visitor.py | 5 +++-- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/pyinfra/test/unit_tests/conftest.py b/pyinfra/test/unit_tests/conftest.py index 1d72983..be06bf2 100644 --- a/pyinfra/test/unit_tests/conftest.py +++ b/pyinfra/test/unit_tests/conftest.py @@ -106,9 +106,9 @@ def analysis_callback(callback): @pytest.fixture -def response_strategy(response_strategy_name): +def response_strategy(response_strategy_name, storage): if response_strategy_name == "storage": - return StorageStrategy() + return StorageStrategy(storage) if response_strategy_name == "forwarding": return ForwardingStrategy() diff --git a/pyinfra/test/unit_tests/queue_visitor_test.py b/pyinfra/test/unit_tests/queue_visitor_test.py index 4fb0b9d..c781429 100644 --- a/pyinfra/test/unit_tests/queue_visitor_test.py +++ b/pyinfra/test/unit_tests/queue_visitor_test.py @@ -3,7 +3,7 @@ import json import pytest -from pyinfra.visitor import get_object_descriptor, StorageStrategy, get_response_object_descriptor +from pyinfra.visitor import get_object_descriptor, get_response_object_descriptor @pytest.fixture() @@ -12,8 +12,8 @@ def body(): @pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session") -@pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session") class TestVisitor: + @pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session") def test_given_a_input_queue_message_callback_pulls_the_data_from_storage( self, visitor, body, storage, bucket_name ): @@ -22,17 +22,15 @@ class TestVisitor: data_received = visitor.load_data(body) assert b"content" == 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): storage.clear_bucket(bucket_name) storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode())) response_body = visitor.load_and_process(body) assert response_body["data"] == "22" + @pytest.mark.parametrize("response_strategy_name", ["storage"], scope="session") def test_visitor_puts_response_on_storage(self, visitor, body, storage, bucket_name): - - visitor.response_strategy = StorageStrategy() - visitor.response_strategy.storage = storage - storage.clear_bucket(bucket_name) storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode())) response_body = visitor(body) diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index 1bff413..1d98eee 100644 --- a/pyinfra/visitor.py +++ b/pyinfra/visitor.py @@ -40,6 +40,9 @@ class ResponseStrategy(abc.ABC): class StorageStrategy(ResponseStrategy): + def __init__(self, storage): + self.storage = storage + def handle_response(self, body): self.storage.put_object(**get_response_object_descriptor(body), data=gzip.compress(json.dumps(body).encode())) body.pop("data") @@ -56,8 +59,6 @@ class QueueVisitor: self.storage = storage self.callback = callback self.response_strategy = response_strategy - # TODO: a little dirty - self.response_strategy.storage = self.storage def load_data(self, body): def download():