Updated unit-tests

Author: codepr

tasq/__init__.py

@@ -4,7 +4,11 @@
 from tasq.worker.jobqueue import JobQueue
 from tasq.worker.executor import ProcessQueueExecutor
 from tasq.remote.client import Client
-from tasq.remote.backend import RedisStoreBackend, RedisBackend
+from tasq.remote.backend import (
+    RedisStoreBackend,
+    RedisBackend,
+    RabbitMQBackend,
+)
 from tasq.remote.runner import Runners
 from tasq.worker.actors import ClientWorker
 from tasq.actors.routers import RoundRobinRouter, actor_pool
@@ -61,6 +65,8 @@ def queue(backend="zmq://localhost:9000", store=None, signkey=None):
         client = Client(_backend)
     elif isinstance(backend, RedisBackend):
         client = Client(BackendConnection(backend))
+    elif isinstance(backend, RabbitMQBackend):
+        client = Client(BackendConnection(backend))
     if store:
         urlstore = urlparse(store)
         assert urlstore.scheme in {

tasq/remote/backend.py

@@ -364,10 +364,10 @@ class RabbitMQBackend:
 
     """Simple Queue with RabbitMQ Backend"""
 
-    def __init__(self, host, port, role, name, namespace="queue"):
+    def __init__(self, amqp_factory, role, name, namespace="queue"):
         """The default connection parameters are: host='localhost', port=5672
         """
-        self._host, self._port = host, port
+        self._driver = amqp_factory()
         assert role in {"receiver", "sender"}, f"Unknown role {role}"
         self._role = role
         self._queue_name = f"{namespace}:{name}"
@@ -378,9 +378,7 @@ def __init__(self, host, port, role, name, namespace="queue"):
         threading.Thread(target=self._start, daemon=True).start()
 
     def _get_channel(self):
-        channel = pika.BlockingConnection(
-            pika.ConnectionParameters(host=self._host, port=self._port)
-        ).channel()
+        channel = self._driver.channel()
         return channel
 
     def _get_job(self, ch, method, _, body):
@@ -408,8 +406,8 @@ def _start(self):
 
     def __repr__(self):
         return (
-            f"RabbitMQBackend(amqp://{self._host}:{self._port}/"
-            f"?name={self._name}, role=self._role)"
+            f"RabbitMQBackend({self._driver}?"
+            f"name={self._name}, role=self._role)"
         )
 
     def put_job(self, serialized_job):
@@ -433,7 +431,8 @@ def get_available_result(self, timeout=None):
             return None
 
     def stop(self):
-        pass
+        ch = self._get_channel()
+        ch.stop()
 
     def close(self):
-        pass
+        self.stop()

tasq/remote/connection.py

@@ -5,14 +5,9 @@
 This module contains classes to define connections using zmq sockets.
 """
 
+import sys
 import zmq
 from urllib.parse import urlparse
-
-try:
-    import redis
-except ImportError:
-    print("You need to install redis python driver to use redis backend")
-
 from .backend import RedisBackend, RabbitMQBackend
 from .sockets import CloudPickleContext, BackendSocket
 from ..exception import BackendCommunicationErrorException
@@ -167,19 +162,42 @@ def from_url(cls, url, signkey=None):
             "port": u.port or 6379,
         }
         if scheme == "redis":
+            try:
+                import redis
+            except ImportError:
+                print(
+                    "You need to install redis python driver to use redis backend"
+                )
+                sys.exit(1)
             conn_args["db"] = int(extraparams.get("db", 0))
             backend = RedisBackend(
                 lambda: redis.StrictRedis(**conn_args), name=name
             )
         else:
-            conn_args["role"] = extraparams.get("role", "sender")
-            backend = RabbitMQBackend(**conn_args)
+            try:
+                import pika
+            except ImportError:
+                print("You need to install pika to use rabbitmq backend")
+                sys.exit(1)
+            role = extraparams.get("role", "sender")
+            backend = RabbitMQBackend(
+                lambda: pika.BlockingConnection(
+                    pika.ConnectionParameters(**conn_args)
+                ),
+                role=role,
+                name=name,
+            )
         return cls(backend, signkey)
 
 
 def connect_redis_backend(
     host, port, db, name, namespace="queue", signkey=None
 ):
+    try:
+        import redis
+    except ImportError:
+        print("You need to install redis python driver to use redis backend")
+        sys.exit(1)
     return BackendConnection(
         RedisBackend(
             lambda: redis.StrictRedis(host, port, db), name, namespace
@@ -191,6 +209,19 @@ def connect_redis_backend(
 def connect_rabbitmq_backend(
     host, port, role, name, namespace="queue", signkey=None
 ):
+    try:
+        import pika
+    except ImportError:
+        print("You need to install pika to use rabbitmq backend")
+        sys.exit(1)
     return BackendConnection(
-        RabbitMQBackend(host, port, role, name, namespace), signkey=signkey,
+        RabbitMQBackend(
+            lambda: pika.BlockingConnection(
+                pika.ConnectionParameters(host=host, port=port)
+            ),
+            role,
+            name,
+            namespace,
+        ),
+        signkey=signkey,
     )

tests/backend_test.py

@@ -1,6 +1,8 @@
 import time
+import queue
 import asyncio
 import unittest
+import threading
 import collections
 from unittest.mock import patch
 from tasq.remote.backend import ZMQBackend, RedisBackend, RabbitMQBackend
@@ -70,6 +72,53 @@ def lrange(self, queue_name, start, end):
         return self.queues[queue_name][start:end]
 
 
+class AMQPMethod:
+    def __init__(self):
+        self.delivery_tag = None
+
+
+class FakeAMQPChannel:
+    def __init__(self):
+        self.run = threading.Event()
+        self.queues = collections.defaultdict(queue.Queue)
+        self.consume_queue = None
+        self.on_message = None
+
+    def basic_publish(self, exchange, queue_name, item):
+        self.queues[queue_name].put(item)
+
+    def basic_consume(self, queue, on_message_callback):
+        self.on_message = on_message_callback
+
+    def basic_ack(self, delivery_tag):
+        pass
+
+    def basic_qos(self, prefetch_count):
+        pass
+
+    def queue_declare(self, queue, durable=True):
+        self.consume_queue = queue
+
+    def start_consuming(self):
+        while not self.run.is_set():
+            item = self.queues[self.consume_queue].get()
+            if not item:
+                break
+            self.on_message(self, AMQPMethod(), None, item)
+
+    def stop(self):
+        self.queues[self.consume_queue].put(None)
+        self.run.set()
+
+
+_amqp_channel = FakeAMQPChannel()
+
+
+class FakeAMQPClient:
+    def channel(self):
+        return _amqp_channel
+
+
 class TestZMQBackend(unittest.TestCase):
     def test_init_zmqbackend(self):
         with patch(
@@ -138,4 +187,10 @@ def test_redis_get_pending_jobs(self):
 
 
 class TestRabbitMQBackend(unittest.TestCase):
-    pass
+    def test_amqp_put_job(self):
+        backend = RabbitMQBackend(
+            lambda: FakeAMQPClient(), role="receiver", name="test-queue"
+        )
+        backend.put_job("test-job")
+        self.assertEqual(backend.get_next_job(), "test-job")
+        backend.stop()

tests/queue_test.py

@@ -141,4 +141,4 @@ def test_multiqueue_put_blocking(self):
         self.assertEqual(tq.pending_jobs(), [])
         self.assertEqual(len(tq.results()), 3)
         self.assertTrue(res)
-        self.assertAlmostEqual(t2 - t1, .1, delta=0.1)
+        self.assertAlmostEqual(t2 - t1, 0.1, delta=0.1)

tests/runner_test.py

@@ -1,9 +1,11 @@
 import queue
 import unittest
 import threading
+from unittest.mock import patch
 from concurrent.futures import Future
 from tasq.remote.runner import Runner, ZMQRunner
 from tasq.job import Job
+from .backend_test import FakeRedisClient, FakeAMQPClient
 
 
 class FakeWorker:
@@ -40,43 +42,55 @@ def test_runner_running(self):
     def test_runner_factory(self):
         from tasq.remote.runner import runner_factory
 
-        zmqrunner = runner_factory.create(
-            "ZMQ_ACTOR_RUNNER", host="localhost", channel=(20000, 20001)
-        )
-        zmqqrunner = runner_factory.create(
-            "ZMQ_QUEUE_RUNNER", host="localhost", channel=(20000, 20001)
-        )
-        redis_actor_runner = runner_factory.create(
-            "REDIS_ACTOR_RUNNER",
-            host="1.2.3.4",
-            port=6333,
-            db=1,
-            name="test-queue",
-        )
-        redis_queue_runner = runner_factory.create(
-            "REDIS_QUEUE_RUNNER",
-            host="1.2.3.4",
-            port=6333,
-            db=1,
-            name="test-queue",
-        )
-        rabbitmq_queue_runner = runner_factory.create(
-            "AMQP_QUEUE_RUNNER",
-            host="1.2.3.4",
-            port=6333,
-            role="sender",
-            name="test-queue",
-        )
-        rabbitmq_actor_runner = runner_factory.create(
-            "AMQP_ACTOR_RUNNER",
-            host="1.2.3.4",
-            port=6333,
-            role="sender",
-            name="test-queue",
-        )
-        self.assertTrue(isinstance(zmqrunner, ZMQRunner))
-        self.assertTrue(isinstance(zmqqrunner, ZMQRunner))
-        self.assertTrue(isinstance(redis_actor_runner, Runner))
-        self.assertTrue(isinstance(redis_queue_runner, Runner))
-        self.assertTrue(isinstance(rabbitmq_actor_runner, Runner))
-        self.assertTrue(isinstance(rabbitmq_queue_runner, Runner))
+        with patch("tasq.remote.runner.ZMQBackend") as zmqmock, patch(
+            "tasq.remote.runner.asyncio.get_event_loop"
+        ) as loopmock, patch("redis.StrictRedis") as rmock, patch(
+            "pika.BlockingConnection"
+        ) as pmock, patch(
+            "tasq.remote.runner.worker.build_jobqueue"
+        ) as jq:
+            jq.return_value = None
+            pmock.return_value = FakeAMQPClient()
+            rmock.return_value = FakeRedisClient()
+            loopmock.return_value = None
+            zmqmock.return_value = None
+            zmqrunner = runner_factory.create(
+                "ZMQ_ACTOR_RUNNER", host="localhost", channel=(20000, 20001)
+            )
+            zmqqrunner = runner_factory.create(
+                "ZMQ_QUEUE_RUNNER", host="localhost", channel=(20000, 20001)
+            )
+            redis_actor_runner = runner_factory.create(
+                "REDIS_ACTOR_RUNNER",
+                host="1.2.3.4",
+                port=6333,
+                db=1,
+                name="test-queue",
+            )
+            redis_queue_runner = runner_factory.create(
+                "REDIS_QUEUE_RUNNER",
+                host="1.2.3.4",
+                port=6333,
+                db=1,
+                name="test-queue",
+            )
+            rabbitmq_queue_runner = runner_factory.create(
+                "AMQP_QUEUE_RUNNER",
+                host="1.2.3.4",
+                port=6333,
+                role="sender",
+                name="test-queue",
+            )
+            rabbitmq_actor_runner = runner_factory.create(
+                "AMQP_ACTOR_RUNNER",
+                host="1.2.3.4",
+                port=6333,
+                role="sender",
+                name="test-queue",
+            )
+            self.assertTrue(isinstance(zmqrunner, ZMQRunner))
+            self.assertTrue(isinstance(zmqqrunner, ZMQRunner))
+            self.assertTrue(isinstance(redis_actor_runner, Runner))
+            self.assertTrue(isinstance(redis_queue_runner, Runner))
+            self.assertTrue(isinstance(rabbitmq_actor_runner, Runner))
+            self.assertTrue(isinstance(rabbitmq_queue_runner, Runner))