diff --git a/pyinfra/test/unit_tests/conftest.py b/pyinfra/test/unit_tests/conftest.py index 0d89c54..328a76a 100644 --- a/pyinfra/test/unit_tests/conftest.py +++ b/pyinfra/test/unit_tests/conftest.py @@ -90,4 +90,11 @@ def queue_manager(queue_manager_name) -> QueueManager: @pytest.fixture def callback(): - return lambda x: x**2 + return lambda x: x * 2 + + +@pytest.fixture +def analysis_callback(callback): + def inner(data: bytes): + return callback(data.decode()) + return inner diff --git a/pyinfra/test/unit_tests/consumer_test.py b/pyinfra/test/unit_tests/consumer_test.py index ffbede8..66e2bb3 100644 --- a/pyinfra/test/unit_tests/consumer_test.py +++ b/pyinfra/test/unit_tests/consumer_test.py @@ -21,7 +21,7 @@ class TestConsumer: ): def produce_items(): - return range(3) + return map(str, range(3)) queue_manager.clear() @@ -33,4 +33,4 @@ class TestConsumer: for _, r in zip(produce_items(), requests): queue_manager.publish_response(r, callback) - assert queue_manager.output_queue.to_list() == [0, 1, 4] + assert queue_manager.output_queue.to_list() == ["00", "11", "22"] diff --git a/pyinfra/test/unit_tests/queue_visitor_test.py b/pyinfra/test/unit_tests/queue_visitor_test.py index ada1488..8f1aba3 100644 --- a/pyinfra/test/unit_tests/queue_visitor_test.py +++ b/pyinfra/test/unit_tests/queue_visitor_test.py @@ -4,8 +4,8 @@ from pyinfra.visitor import QueueVisitor, get_object_descriptor @pytest.fixture() -def visitor(storage, callback): - return QueueVisitor(storage, callback) +def visitor(storage, analysis_callback): + return QueueVisitor(storage, analysis_callback) @pytest.fixture() @@ -21,3 +21,9 @@ class TestStorage: storage.put_object(**get_object_descriptor(body), data=b"content") data_received = visitor.load_data(body) assert b"content" == data_received + + def test_visitor_call_method_pulls_data_and_executes_callback(self, visitor, body, storage, bucket_name): + storage.clear_bucket(bucket_name) + storage.put_object(**get_object_descriptor(body), data="2".encode()) + response_body = visitor(body) + assert response_body == "22" diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index fbe901d..a011d09 100644 --- a/pyinfra/visitor.py +++ b/pyinfra/visitor.py @@ -28,8 +28,8 @@ class QueueVisitor: def load_data(self, body): return self.storage.get_object(**get_object_descriptor(body)) - def process_data(self): - pass + def process_data(self, data): + return self.callback(data) def __call__(self, body): - pass + return self.process_data(self.load_data(body))