from __future__ import absolute_import import socket import types from anyjson import dumps, loads from collections import defaultdict from contextlib import contextmanager from itertools import count from kombu import Connection, Exchange, Queue, Consumer, Producer from kombu.exceptions import InconsistencyError, VersionMismatch from kombu.five import Empty, Queue as _Queue from kombu.transport import virtual from kombu.utils import eventio # patch poll from kombu.tests.case import ( Case, ContextMock, Mock, call, module_exists, skip_if_not_module, patch, ) class JSONEqual(object): # The order in which a dict is serialized to json depends on the hashseed # so we have this to support json in .assert_has_call*. def __init__(self, expected): self.expected = expected def __eq__(self, other): return loads(other) == loads(self.expected) def __str__(self): return self.expected def __repr__(self): return '(json)%r' % (self.expected,) class _poll(eventio._select): def register(self, fd, flags): if flags & eventio.READ: self._rfd.add(fd) def poll(self, timeout): events = [] for fd in self._rfd: if fd.data: events.append((fd.fileno(), eventio.READ)) return events eventio.poll = _poll # must import after poller patch from kombu.transport import redis # noqa class ResponseError(Exception): pass class Client(object): queues = {} sets = defaultdict(set) hashes = defaultdict(dict) shard_hint = None def __init__(self, db=None, port=None, connection_pool=None, **kwargs): self._called = [] self._connection = None self.bgsave_raises_ResponseError = False self.connection = self._sconnection(self) def bgsave(self): self._called.append('BGSAVE') if self.bgsave_raises_ResponseError: raise ResponseError() def delete(self, key): self.queues.pop(key, None) def exists(self, key): return key in self.queues or key in self.sets def hset(self, key, k, v): self.hashes[key][k] = v def hget(self, key, k): return self.hashes[key].get(k) def hdel(self, key, k): self.hashes[key].pop(k, None) def sadd(self, key, member, *args): self.sets[key].add(member) zadd = sadd def smembers(self, key): return self.sets.get(key, set()) def srem(self, key, *args): self.sets.pop(key, None) zrem = srem def llen(self, key): try: return self.queues[key].qsize() except KeyError: return 0 def lpush(self, key, value): self.queues[key].put_nowait(value) def parse_response(self, connection, type, **options): cmd, queues = self.connection._sock.data.pop() assert cmd == type self.connection._sock.data = [] if type == 'BRPOP': item = self.brpop(queues, 0.001) if item: return item raise Empty() def brpop(self, keys, timeout=None): key = keys[0] try: item = self.queues[key].get(timeout=timeout) except Empty: pass else: return key, item def rpop(self, key): try: return self.queues[key].get_nowait() except KeyError: pass def __contains__(self, k): return k in self._called def pipeline(self): return Pipeline(self) def encode(self, value): return str(value) def _new_queue(self, key): self.queues[key] = _Queue() class _sconnection(object): disconnected = False class _socket(object): blocking = True filenos = count(30) def __init__(self, *args): self._fileno = next(self.filenos) self.data = [] def fileno(self): return self._fileno def setblocking(self, blocking): self.blocking = blocking def __init__(self, client): self.client = client self._sock = self._socket() def disconnect(self): self.disconnected = True def send_command(self, cmd, *args): self._sock.data.append((cmd, args)) def info(self): return {'foo': 1} def pubsub(self, *args, **kwargs): connection = self.connection class ConnectionPool(object): def get_connection(self, *args, **kwargs): return connection self.connection_pool = ConnectionPool() return self def __repr__(self): return '