pyinfra/test/unit_tests/queue_visitor_test.py

import json

import pytest
from funcy import first

from pyinfra.utils.encoding import decompress
from test.utils.storage import pack_for_upload


@pytest.fixture()
def body():
    return {"dossierId": "dossier_id", "fileId": "file_id"}


@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], 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, file_descriptor_manager
    ):
        storage.clear_bucket(bucket_name)
        storage.put_object(**file_descriptor_manager.get_input_object_descriptor(body), data=pack_for_upload(b"content"))
        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, file_descriptor_manager):
        storage.clear_bucket(bucket_name)
        storage.put_object(**file_descriptor_manager.get_input_object_descriptor(body), data=pack_for_upload(b"2"))
        response_body = visitor.load_items_from_storage_and_process_with_callback(body)
        assert response_body["analysis_payloads"] == ["22"]

    @pytest.mark.parametrize("response_strategy_name", ["storage"], scope="session")
    def test_visitor_puts_response_on_storage(self, visitor, body, storage, bucket_name, file_descriptor_manager):
        storage.clear_bucket(bucket_name)
        storage.put_object(**file_descriptor_manager.get_input_object_descriptor(body), data=pack_for_upload(b"2"))
        response_body = visitor(body)
        assert "data" not in response_body
        assert json.loads(
            decompress(storage.get_object(bucket_name=bucket_name, object_name=first(response_body["response_files"])))
        )["analysis_payloads"] == ["22"]