Modified patch_connection to allow Redis connection to be easily mocked.

main
Selwin Ong 9 years ago
parent e6a15c57b3
commit 0e4112da7a

@ -18,12 +18,14 @@ def fix_return_type(func):
return _inner return _inner
PATCHED_METHODS = ['_setex', '_lrem', '_zadd', '_pipeline', '_ttl']
def patch_connection(connection): def patch_connection(connection):
if not isinstance(connection, StrictRedis): if not isinstance(connection, StrictRedis):
raise ValueError('A StrictRedis or Redis connection is required.') raise ValueError('A StrictRedis or Redis connection is required.')
# Don't patch already patches objects # Don't patch already patches objects
PATCHED_METHODS = ['_setex', '_lrem', '_zadd', '_pipeline', '_ttl']
if all([hasattr(connection, attr) for attr in PATCHED_METHODS]): if all([hasattr(connection, attr) for attr in PATCHED_METHODS]):
return connection return connection
@ -35,6 +37,7 @@ def patch_connection(connection):
connection._ttl = fix_return_type(partial(StrictRedis.ttl, connection)) connection._ttl = fix_return_type(partial(StrictRedis.ttl, connection))
if hasattr(connection, 'pttl'): if hasattr(connection, 'pttl'):
connection._pttl = fix_return_type(partial(StrictRedis.pttl, connection)) connection._pttl = fix_return_type(partial(StrictRedis.pttl, connection))
elif isinstance(connection, StrictRedis): elif isinstance(connection, StrictRedis):
connection._setex = connection.setex connection._setex = connection.setex
connection._lrem = connection.lrem connection._lrem = connection.lrem

@ -5,7 +5,7 @@ from __future__ import (absolute_import, division, print_function,
from datetime import datetime from datetime import datetime
import time import time
from tests import RQTestCase from tests import fixtures, RQTestCase
from tests.helpers import strip_microseconds from tests.helpers import strip_microseconds
from rq.compat import PY2, as_text from rq.compat import PY2, as_text
@ -16,8 +16,6 @@ from rq.registry import DeferredJobRegistry
from rq.utils import utcformat from rq.utils import utcformat
from rq.worker import Worker from rq.worker import Worker
from . import fixtures
try: try:
from cPickle import loads, dumps from cPickle import loads, dumps
except ImportError: except ImportError:

Loading…
Cancel
Save