added tests for queue visitor

Merge in RR/pyinfra from prometheus-tunneling to pika_encapsulation

Squashed commit of the following:

commit 350448ec0f14849844deaa1a86ba4397ab3ebf3c
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Wed Mar 9 16:52:02 2022 +0100

    quickfix

commit ee88be4c80abdf597013c743ce2d490ad2b3e029
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Wed Mar 9 15:26:19 2022 +0100

    added prometheus endpoint to tunnel metrics from analysis endpoint

config.yaml

@@ -1,8 +1,8 @@
 service:
-  logging_level: $LOGGING_LEVEL_ROOT|DEBUG  # Logging level for service logger
+  logging_level: $LOGGING_LEVEL_ROOT|INFO  # Logging level for service logger
   response:
     type: $RESPONSE_TYPE|"file"  # Whether the analysis response is stored as file on storage or sent as stream
-    extension: $RESPONSE_FILE_EXTENSION|".IMAGE_INFO.json.gz"  # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz}
+    extension: $RESPONSE_FILE_EXTENSION|"IMAGE_INFO.json.gz"  # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz}
     key: $RESPONSE_KEY|"imageMetadata" # the key of the result {result, imageMetadata}
 
 probing_webserver:
@@ -37,7 +37,8 @@ storage:
 
   backend: $STORAGE_BACKEND|s3  # The type of storage to use {s3, azure}
   bucket: $STORAGE_BUCKET|"pyinfra-test-bucket"  # The bucket / container to pull files specified in queue requests from
-  target_file_extension: $TARGET_FILE_EXTENSION|".ORIGIN.pdf.gz"  # {.TEXT.json.gz | .ORIGIN.pdf.gz} Defines type of file to pull from storage
+  # TODO: Caller should specify exact file name, including extension!
+  target_file_extension: $TARGET_FILE_EXTENSION|"ORIGIN.pdf.gz"  # {.TEXT.json.gz | .ORIGIN.pdf.gz} Defines type of file to pull from storage
 
   s3:
     endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000"

pyinfra/config.py

@@ -1,11 +1,23 @@
 """Implements a config object with dot-indexing syntax."""
 
-
 from envyaml import EnvYAML
 
 from pyinfra.locations import CONFIG_FILE
 
 
+def make_art():
+    return """
+______     _____       __           
+| ___ \   |_   _|     / _|          
+| |_/ /   _ | | _ __ | |_ _ __ __ _ 
+|  __/ | | || || '_ \|  _| '__/ _` |
+| |  | |_| || || | | | | | | | (_| |
+\_|   \__, \___/_| |_|_| |_|  \__,_|
+       __/ |                        
+      |___/                         
+"""
+
+
 def _get_item_and_maybe_make_dotindexable(container, item):
     ret = container[item]
     return DotIndexable(ret) if isinstance(ret, dict) else ret

pyinfra/flask.py

@@ -1,3 +1,4 @@
+import requests
 from flask import Flask, jsonify
 from waitress import serve
 
@@ -5,7 +6,6 @@ from pyinfra.config import CONFIG
 
 
 def run_probing_webserver(app, host=None, port=None, mode=None):
-
     if not host:
         host = CONFIG.probing_webserver.host
 
@@ -23,7 +23,6 @@ def run_probing_webserver(app, host=None, port=None, mode=None):
 
 
 def set_up_probing_webserver():
-
     app = Flask(__name__)
 
     @app.route("/ready", methods=["GET"])
@@ -38,4 +37,10 @@ def set_up_probing_webserver():
         resp.status_code = 200
         return resp
 
+    @app.route("/prometheus", methods=["GET"])
+    def get_analysis_prometheus_endpoint():
+        prom_endpoint = f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus"
+        metric = requests.get(prom_endpoint)
+        return metric.text
+
     return app

pyinfra/pyampq/connection/connection.py

@@ -0,0 +1,7 @@
+import abc
+
+
+class Connection(abc.ABC):
+    @abc.abstractmethod
+    def establish(self):
+        pass

pyinfra/pyampq/connection/mock_connection.py

@@ -0,0 +1,6 @@
+from pyampq.connection.connection import Connection
+
+
+class MockConnection(Connection):
+    def establish(self):
+        pass

pyinfra/pyampq/connection/pika_connection.py

@@ -0,0 +1,3 @@
+class PikaConnection:
+    def establish(self):
+        pass

pyinfra/pyampq/consumer.py

@@ -0,0 +1,13 @@
+from pyinfra.pyampq.queue_manager.queue_manager import QueueManager
+
+
+class Consumer:
+    def __init__(self, callback, queue_manager: QueueManager):
+        self.queue_manager = queue_manager
+        self.callback = callback
+
+    def consume_and_publish(self):
+        yield from self.queue_manager.consume_and_publish(self.callback)
+
+    def consume(self):
+        yield from self.queue_manager.consume()

pyinfra/pyampq/queue_manager/pika_queue_manager.py

@@ -0,0 +1,98 @@
+import json
+import logging
+
+import pika
+
+from pyinfra.config import CONFIG
+from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle
+
+logger = logging.getLogger("pika")
+logger.setLevel(logging.WARNING)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(CONFIG.service.logging_level)
+
+
+def monkey_patch_queue_handle(channel, queue) -> QueueHandle:
+
+    empty_message = (None, None, None)
+
+    def is_empty_message(message):
+        return message == empty_message
+
+    queue_handle = QueueHandle()
+    queue_handle.empty = lambda: is_empty_message(channel.basic_get(queue))
+
+    def produce_items():
+
+        while True:
+            message = channel.basic_get(queue)
+
+            if is_empty_message(message):
+                break
+
+            method_frame, properties, body = message
+            channel.basic_ack(method_frame.delivery_tag)
+            yield json.loads(body)
+
+    queue_handle.to_list = lambda: list(produce_items())
+
+    return queue_handle
+
+
+def get_connection():
+
+    credentials = pika.PlainCredentials(username=CONFIG.rabbitmq.user, password=CONFIG.rabbitmq.password)
+
+    kwargs = {"host": CONFIG.rabbitmq.host, "port": CONFIG.rabbitmq.port, "credentials": credentials}
+
+    parameters = pika.ConnectionParameters(**kwargs)
+
+    connection = pika.BlockingConnection(parameters=parameters)
+
+    return connection
+
+
+class PikaQueueManager(QueueManager):
+    def __init__(self, input_queue, output_queue):
+        super().__init__(input_queue, output_queue)
+        connection = get_connection()
+        self.channel = connection.channel()
+        self.channel.queue_declare(input_queue)
+        self.channel.queue_declare(output_queue)
+
+    def publish_request(self, request):
+        self.channel.basic_publish("", self._input_queue, json.dumps(request).encode())
+
+    def publish_response(self, message, callback):
+
+        frame, properties, body = message
+        response = json.dumps(callback(json.loads(body))).encode()
+
+        self.channel.basic_publish("", self._output_queue, response)
+        self.channel.basic_ack(frame.delivery_tag)
+
+    def pull_request(self):
+        return self.channel.basic_get(self._input_queue)
+
+    def consume(self):
+        return self.channel.consume(self._input_queue)
+
+    def consume_and_publish(self, callback):
+
+        logger.info(f"Consuming with callback {callback.__name__}")
+
+        for message in self.consume():
+            self.publish_response(message, callback)
+
+    def clear(self):
+        self.channel.queue_purge(self._input_queue)
+        self.channel.queue_purge(self._output_queue)
+
+    @property
+    def input_queue(self) -> QueueHandle:
+        return monkey_patch_queue_handle(self.channel, self._input_queue)
+
+    @property
+    def output_queue(self) -> QueueHandle:
+        return monkey_patch_queue_handle(self.channel, self._output_queue)

pyinfra/pyampq/queue_manager/queue_manager.py

@@ -0,0 +1,47 @@
+import abc
+
+
+class QueueHandle:
+    def empty(self) -> bool:
+        raise NotImplemented()
+
+    def to_list(self) -> list:
+        raise NotImplemented()
+
+
+class QueueManager(abc.ABC):
+    def __init__(self, input_queue, output_queue):
+        self._input_queue = input_queue
+        self._output_queue = output_queue
+
+    @abc.abstractmethod
+    def publish_request(self, request):
+        pass
+
+    @abc.abstractmethod
+    def publish_response(self, response, callback):
+        pass
+
+    @abc.abstractmethod
+    def pull_request(self):
+        pass
+
+    @abc.abstractmethod
+    def consume(self):
+        pass
+
+    @abc.abstractmethod
+    def clear(self):
+        pass
+
+    @abc.abstractmethod
+    def input_queue(self) -> QueueHandle:
+        pass
+
+    @abc.abstractmethod
+    def output_queue(self) -> QueueHandle:
+        pass
+
+    @abc.abstractmethod
+    def consume_and_publish(self, callback):
+        pass

pyinfra/test/queue_manager_mock.py

@@ -0,0 +1,43 @@
+from pyinfra.pyampq.queue_manager.queue_manager import QueueManager, QueueHandle
+from pyinfra.test.queue_mock import QueueMock
+
+
+def monkey_patch_queue_handle(queue) -> QueueHandle:
+    queue_handle = QueueHandle()
+    queue_handle.empty = lambda: not queue
+    queue_handle.to_list = lambda: list(queue)
+    return queue_handle
+
+
+class QueueManagerMock(QueueManager):
+    def __init__(self, input_queue, output_queue):
+        super().__init__(QueueMock(), QueueMock())
+
+    def publish_request(self, request):
+        self._input_queue.append(request)
+
+    def publish_response(self, message, callback):
+        self._output_queue.append(callback(message))
+
+    def pull_request(self):
+        return self._input_queue.popleft()
+
+    def consume(self):
+        while self._input_queue:
+            yield self.pull_request()
+
+    def consume_and_publish(self, callback):
+        for message in self.consume():
+            self.publish_response(message, callback)
+
+    def clear(self):
+        self._input_queue.clear()
+        self._output_queue.clear()
+
+    @property
+    def input_queue(self) -> QueueHandle:
+        return monkey_patch_queue_handle(self._input_queue)
+
+    @property
+    def output_queue(self) -> QueueHandle:
+        return monkey_patch_queue_handle(self._output_queue)

pyinfra/test/queue_mock.py

@@ -0,0 +1,5 @@
+from collections import deque
+
+
+class QueueMock(deque):
+    pass

pyinfra/test/unit_tests/conftest.py

@@ -6,12 +6,15 @@ import pytest
 
 from pyinfra.exceptions import UnknownClient
 from pyinfra.locations import TEST_DIR
+from pyinfra.pyampq.queue_manager.pika_queue_manager import PikaQueueManager
+from pyinfra.pyampq.queue_manager.queue_manager import QueueManager
 from pyinfra.storage.adapters.azure import AzureStorageAdapter
 from pyinfra.storage.adapters.s3 import S3StorageAdapter
 from pyinfra.storage.clients.azure import get_azure_client
 from pyinfra.storage.clients.s3 import get_s3_client
 from pyinfra.storage.storage import Storage
 from pyinfra.test.config import CONFIG
+from pyinfra.test.queue_manager_mock import QueueManagerMock
 from pyinfra.test.storage.adapter_mock import StorageAdapterMock
 from pyinfra.test.storage.client_mock import StorageClientMock
 
@@ -19,7 +22,7 @@ logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def bucket_name():
     return "pyinfra-test-bucket"
 
@@ -51,7 +54,7 @@ def mock_make_load_data():
     return load_data
 
 
-@pytest.fixture(params=["minio", "aws"])
+@pytest.fixture(params=["minio", "aws"], scope="session")
 def storage(client_name, bucket_name, request):
     logger.debug("Setup for storage")
     storage = Storage(get_adapter(client_name, request.param))
@@ -71,3 +74,20 @@ def get_adapter(client_name, s3_backend):
         return S3StorageAdapter(get_s3_client(CONFIG.storage[s3_backend]))
     else:
         raise UnknownClient(client_name)
+
+
+def get_queue_manager(queue_manager_name) -> QueueManager:
+    if queue_manager_name == "mock":
+        return QueueManagerMock("input", "output")
+    if queue_manager_name == "pika":
+        return PikaQueueManager("input", "output")
+
+
+@pytest.fixture(scope="session")
+def queue_manager(queue_manager_name) -> QueueManager:
+    return get_queue_manager(queue_manager_name)
+
+
+@pytest.fixture
+def callback():
+    return lambda x: x**2

pyinfra/test/unit_tests/consumer_test.py

@@ -0,0 +1,36 @@
+import pytest
+
+from pyinfra.pyampq.consumer import Consumer
+
+
+@pytest.fixture
+def consumer(queue_manager, callback):
+    return Consumer(callback, queue_manager)
+
+
+@pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session")
+class TestConsumer:
+
+    def test_consuming_empty_input_queue_does_not_put_anything_on_output_queue(self, consumer, queue_manager):
+        queue_manager.clear()
+        consumer.consume()
+        assert queue_manager.output_queue.empty()
+
+    def test_consuming_nonempty_input_queue_puts_messages_on_output_queue_in_fifo_order(
+            self, consumer, queue_manager, callback
+    ):
+
+        def produce_items():
+            return range(3)
+
+        queue_manager.clear()
+
+        for item in produce_items():
+            queue_manager.publish_request(item)
+
+        requests = consumer.consume()
+
+        for _, r in zip(produce_items(), requests):
+            queue_manager.publish_response(r, callback)
+
+        assert queue_manager.output_queue.to_list() == [0, 1, 4]

pyinfra/test/unit_tests/queue_visitor_test.py

@@ -0,0 +1,23 @@
+import pytest
+
+from pyinfra.visitor import QueueVisitor, get_object_descriptor
+
+
+@pytest.fixture()
+def visitor(storage, callback):
+    return QueueVisitor(storage, callback)
+
+
+@pytest.fixture()
+def body():
+    return {"dossierId": "folder", "fileId": "file"}
+
+
+@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
+class TestStorage:
+
+    def test_given_a_input_queue_message_callback_pulls_the_data_from_storage(self, visitor, body, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
+        storage.put_object(**get_object_descriptor(body), data=b"content")
+        data_received = visitor.load_data(body)
+        assert b"content" == data_received

pyinfra/test/unit_tests/storage_test.py

@@ -9,7 +9,7 @@ logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"])
+@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
 class TestStorage:
     def test_clearing_bucket_yields_empty_bucket(self, storage, bucket_name):
         storage.clear_bucket(bucket_name)
@@ -17,26 +17,31 @@ class TestStorage:
         assert not {*data_received}
 
     def test_getting_object_put_in_bucket_is_object(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "file", b"content")
         data_received = storage.get_object(bucket_name, "file")
         assert b"content" == data_received
 
     def test_getting_nested_object_put_in_bucket_is_nested_object(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "folder/file", b"content")
         data_received = storage.get_object(bucket_name, "folder/file")
         assert b"content" == data_received
 
     def test_getting_objects_put_in_bucket_are_objects(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "file1", b"content 1")
         storage.put_object(bucket_name, "folder/file2", b"content 2")
         data_received = storage.get_all_objects(bucket_name)
         assert {b"content 1", b"content 2"} == {*data_received}
 
     def test_make_bucket_produces_bucket(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.make_bucket(bucket_name)
         assert storage.has_bucket(bucket_name)
 
     def test_listing_bucket_files_yields_all_files_in_bucket(self, storage, bucket_name):
+        storage.clear_bucket(bucket_name)
         storage.put_object(bucket_name, "file1", b"content 1")
         storage.put_object(bucket_name, "file2", b"content 2")
         full_names_received = storage.get_all_object_names(bucket_name)

pyinfra/visitor.py

@@ -0,0 +1,35 @@
+from operator import itemgetter
+from typing import Callable
+
+from pyinfra.config import CONFIG
+from pyinfra.storage.storage import Storage
+
+
+def get_object_name(body):
+    dossier_id, file_id = itemgetter("dossierId", "fileId")(body)
+    object_name = f"{dossier_id}/{file_id}.{CONFIG.target_file_extension}"
+    return object_name
+
+
+def get_object_descriptor(body):
+    return {
+        "bucket_name": CONFIG.storage.bucket,
+        # TODO: send complete filename in request, as discussed!
+        "object_name": get_object_name(body)
+    }
+
+
+class QueueVisitor:
+
+    def __init__(self, storage: Storage, callback: Callable):
+        self.storage = storage
+        self.callback = callback
+
+    def load_data(self, body):
+        return self.storage.get_object(**get_object_descriptor(body))
+
+    def process_data(self):
+        pass
+
+    def __call__(self, body):
+        pass

src/serve.py

@@ -8,7 +8,7 @@ from pyinfra.callback import (
     make_retry_callback,
     make_callback_for_output_queue,
 )
-from pyinfra.config import CONFIG
+from pyinfra.config import CONFIG, make_art
 from pyinfra.consume import consume, ConsumerError
 from pyinfra.core import make_payload_processor, make_storage_data_loader, make_analyzer
 from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
@@ -49,6 +49,7 @@ def make_callback():
 
 def main():
     # TODO: implement meaningful checks
+    logging.info(make_art())
     webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
     logging.info("Starting webserver...")
     webserver.start()