Add basic backend testing

Author: codepr

.gitignore

@@ -6,3 +6,4 @@ tags
 *.log
 *.conf
 *.egg-info
+.coverage

tasq/__init__.py

@@ -4,7 +4,7 @@
 from tasq.worker.jobqueue import JobQueue
 from tasq.worker.executor import ProcessQueueExecutor
 from tasq.remote.client import Client
-from tasq.remote.backend import RedisStoreBackend
+from tasq.remote.backend import RedisStoreBackend, RedisBackend
 from tasq.remote.runner import Runners
 from tasq.worker.actors import ClientWorker
 from tasq.actors.routers import RoundRobinRouter, actor_pool
@@ -27,7 +27,7 @@
 }
 
 
-def queue(url="zmq://localhost:9000", store=None, signkey=None):
+def queue(backend="zmq://localhost:9000", store=None, signkey=None):
     """
     Create a TasqQueue instance.
     The formats accepted for the backends are:
@@ -53,17 +53,20 @@ def queue(url="zmq://localhost:9000", store=None, signkey=None):
     :param signkey: A string representing a shared key, sign data with a shared
                     key
     """
-    url_parsed = urlparse(url)
-    scheme = url_parsed.scheme or "zmq"
-    assert scheme in _backends, f"Unsupported {scheme} as backend"
-    backend = _backends[scheme].from_url(url, signkey)
-    client = Client(backend)
+    if isinstance(backend, str):
+        url_parsed = urlparse(backend)
+        scheme = url_parsed.scheme or "zmq"
+        assert scheme in _backends, f"Unsupported {scheme} as backend"
+        _backend = _backends[scheme].from_url(backend, signkey)
+        client = Client(_backend)
+    elif isinstance(backend, RedisBackend):
+        client = Client(BackendConnection(backend))
     if store:
         urlstore = urlparse(store)
         assert urlstore.scheme in {
             "redis"
         }, f"Unknown {urlstore.scheme} as store"
-        db = int(urlstore.path.split("/")[-1]) if url.query else 0
+        db = int(urlstore.path.split("/")[-1]) if urlstore.query else 0
         store = RedisStoreBackend(urlstore.hostname, urlstore.port, db)
     return TasqQueue(client, store)
 

tasq/remote/backend.py

@@ -138,11 +138,11 @@ class RedisQueue:
 
         log = get_logger(__name__)
 
-        def __init__(self, name, host, port, db, namespace="queue"):
+        def __init__(self, name, redis_driver, namespace="queue"):
             """The default connection parameters are: host='localhost',
             port=6379, db=0
             """
-            self._db = redis.StrictRedis(host=host, port=port, db=db)
+            self._db = redis_driver
 
             try:
                 _ = self._db.dbsize()
@@ -152,6 +152,9 @@ def __init__(self, name, host, port, db, namespace="queue"):
             self._queue_name = f"{namespace}:{name}"
             self._work_queue_name = f"{namespace}:{name}:work"
 
+        def __repr__(self):
+            return f"{self._db}:({self._queue_name}, {self._work_queue_name})"
+
         def qsize(self):
             """Return the approximate size of the queue."""
             return self._db.llen(self._queue_name)
@@ -245,22 +248,16 @@ def list_working_items(self):
         def close(self):
             self._db.connection_pool.disconnect()
 
-    def __init__(self, host, port, db, name, namespace="queue"):
+    def __init__(self, redis_factory, name, namespace="queue"):
 
-        self._rq = self.RedisQueue(name, host, port, db, namespace)
+        self._rq = self.RedisQueue(name, redis_factory(), namespace)
         self._rq_res = self.RedisQueue(
-            f"{name}:result", host, port, db, namespace
+            f"{name}:result", redis_factory(), namespace
         )
-        self._host = host
-        self._port = port
-        self._db = db
         self._namespace = f"{namespace}:{name}"
 
     def __repr__(self):
-        return (
-            f"RedisBackend(redis://{self._host}:{self._port}/{self._db}"
-            f"?name={self._namespace}"
-        )
+        return f"RedisBackend(redis://{self._rq}"
 
     def put_job(self, serialized_job):
         self._rq.put(serialized_job)

tasq/remote/connection.py

@@ -7,6 +7,12 @@
 
 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
@@ -159,11 +165,12 @@ def from_url(cls, url, signkey=None):
         conn_args = {
             "host": u.hostname or "localhost",
             "port": u.port or 6379,
-            "name": name,
         }
         if scheme == "redis":
             conn_args["db"] = int(extraparams.get("db", 0))
-            backend = RedisBackend(**conn_args)
+            backend = RedisBackend(
+                lambda: redis.StrictRedis(**conn_args), name=name
+            )
         else:
             conn_args["role"] = extraparams.get("role", "sender")
             backend = RabbitMQBackend(**conn_args)
@@ -174,7 +181,8 @@ def connect_redis_backend(
     host, port, db, name, namespace="queue", signkey=None
 ):
     return BackendConnection(
-        RedisBackend(host, port, db, name, namespace), signkey=signkey,
+        RedisBackend(lambda: redis.StrictRedis(host, port, db), name, namespace),
+        signkey=signkey,
     )
 
 

tasq/remote/runner.py

@@ -44,6 +44,7 @@ def start(self):
         for incoming tasks and run the asyncio loop forever
         """
         self._run = True
+        self._log.debug("Listening on %s", self._backend)
         self.run()
 
     def run(self):

tests/backend_test.py

@@ -0,0 +1,141 @@
+import time
+import asyncio
+import unittest
+import collections
+from unittest.mock import patch
+from tasq.remote.backend import ZMQBackend, RedisBackend, RabbitMQBackend
+
+
+class FakeSocket:
+    def __init__(self):
+        self.bind_url = u""
+        self.data_sent = None
+
+    async def send_data(self, data, flags, signkey):
+        self.data_sent = (data, flags, signkey)
+
+    async def recv_data(self, unpickle, flags, signkey):
+        return self.data_sent
+
+    def bind(self, url):
+        self.bind_url = url
+
+    def close(self):
+        pass
+
+
+class FakeRedisPool:
+    def disconnect(self):
+        pass
+
+
+class FakeRedisClient:
+    def __init__(self):
+        self.queues = collections.defaultdict(list)
+        self.connection_pool = FakeRedisPool()
+
+    def dbsize(self):
+        return 0
+
+    def llen(self, queue_name):
+        return len(self.queues[queue_name])
+
+    def lpush(self, queue_name, item):
+        self.queues[queue_name].append(item)
+
+    def brpop(self, queue_name, timeout=0):
+        if timeout and not self.queues[queue_name]:
+            time.sleep(timeout)
+            return None
+        return (
+            None,
+            self.queues[queue_name],
+        )
+
+    def rpop(self, queue_name):
+        return self.brpop(queue_name)
+
+    def brpoplpush(self, queue_name, queue_name_alt, timeout=0):
+        if timeout and not self.queues[queue_name]:
+            time.sleep(timeout)
+            return None
+        item = self.queues[queue_name]
+        self.queues[queue_name_alt].append(item)
+        return item
+
+    def rpoplpush(self, queue_name, queue_name_alt):
+        return self.brpoplpush(queue_name, queue_name_alt)
+
+    def lrange(self, queue_name, start, end):
+        return self.queues[queue_name][start:end]
+
+
+class TestZMQBackend(unittest.TestCase):
+    def test_init_zmqbackend(self):
+        with patch(
+            "tasq.remote.backend.AsyncCloudPickleContext.socket"
+        ) as mock:
+            mock.side_effect = lambda _: FakeSocket()
+            backend = ZMQBackend("localhost", 10000, 10001)
+            backend.bind()
+            self.assertEqual(
+                backend._pull_socket.bind_url, "tcp://localhost:10001"
+            )
+            self.assertEqual(
+                backend._push_socket.bind_url, "tcp://localhost:10000"
+            )
+            backend.stop()
+
+    def test_send_zmqbackend(self):
+        fake_socket = FakeSocket()
+        with patch(
+            "tasq.remote.backend.AsyncCloudPickleContext.socket"
+        ) as mock:
+            mock.return_value = fake_socket
+            backend = ZMQBackend("localhost", 10000, 10001)
+            backend.bind()
+            asyncio.run(backend.send("hello"))
+            self.assertEqual(fake_socket.data_sent, ("hello", 0, None))
+
+    def test_recv_zmqbackend(self):
+        fake_socket = FakeSocket()
+        with patch(
+            "tasq.remote.backend.AsyncCloudPickleContext.socket"
+        ) as mock:
+            mock.return_value = fake_socket
+            backend = ZMQBackend("localhost", 10000, 10001)
+            backend.bind()
+            asyncio.run(backend.send("hello"))
+            self.assertEqual(fake_socket.data_sent, ("hello", 0, None))
+            payload = asyncio.run(backend.recv())
+            self.assertEqual(payload, ("hello", 0, None))
+
+
+class TestRedisBackend(unittest.TestCase):
+    def test_redis_put_job(self):
+        backend = RedisBackend(lambda: FakeRedisClient(), name="test-queue")
+        backend.put_job("test-job")
+        self.assertEqual(backend.get_next_job(), ["test-job"])
+        backend.close()
+
+    def test_redis_put_result(self):
+        backend = RedisBackend(lambda: FakeRedisClient(), name="test-queue")
+        backend.put_result("test-job-result")
+        self.assertEqual(backend.get_available_result(), ["test-job-result"])
+        backend.close()
+
+    def test_redis_get_next_job(self):
+        backend = RedisBackend(lambda: FakeRedisClient(), name="test-queue")
+        backend.put_job("test-job")
+        self.assertEqual(backend.get_next_job(), ["test-job"])
+        backend.close()
+
+    def test_redis_get_pending_jobs(self):
+        backend = RedisBackend(lambda: FakeRedisClient(), name="test-queue")
+        backend.put_job("test-job")
+        self.assertEqual(backend.get_pending_jobs(), ([], []))
+        backend.close()
+
+
+class TestRabbitMQBackend(unittest.TestCase):
+    pass