test refac

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,5 +1,5 @@
 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}
@@ -37,6 +37,7 @@ 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
+  # 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:

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,99 @@
+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)
+            print(message)
+
+            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/consumer_test.py

@@ -0,0 +1,46 @@
+import pytest
+
+from pyinfra.pyampq.consumer import Consumer
+from pyinfra.pyampq.queue_manager.pika_queue_manager import PikaQueueManager
+from pyinfra.pyampq.queue_manager.queue_manager import QueueManager
+from pyinfra.test.queue_manager_mock import QueueManagerMock
+
+
+@pytest.fixture
+def callback():
+    return lambda x: x**2
+
+
+@pytest.fixture(scope="session")
+def queue_manager() -> QueueManager:
+    return QueueManagerMock("input", "output")
+    return PikaQueueManager("input", "output")
+
+
+@pytest.fixture
+def consumer(queue_manager, callback):
+    return Consumer(callback, queue_manager)
+
+
+def test_consuming_empty_input_queue_does_not_put_anything_on_output_queue(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(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]

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()