127 lines
3.3 KiB
Python
127 lines
3.3 KiB
Python
import json
|
|
from collections import defaultdict
|
|
|
|
|
|
class FakePipeline:
|
|
def __init__(self, client):
|
|
self.client = client
|
|
self.operations = []
|
|
|
|
def __getattr__(self, name):
|
|
def wrapper(*args, **kwargs):
|
|
self.operations.append((name, args, kwargs))
|
|
return self
|
|
|
|
return wrapper
|
|
|
|
def execute(self):
|
|
results = []
|
|
for name, args, kwargs in self.operations:
|
|
results.append(getattr(self.client, name)(*args, **kwargs))
|
|
self.operations.clear()
|
|
return results
|
|
|
|
|
|
class FakePubSub:
|
|
def __init__(self):
|
|
self.channels = []
|
|
self.messages = []
|
|
self.closed = False
|
|
|
|
def subscribe(self, channel):
|
|
self.channels.append(channel)
|
|
|
|
def unsubscribe(self, channel):
|
|
if channel in self.channels:
|
|
self.channels.remove(channel)
|
|
|
|
def get_message(self, timeout=1.0):
|
|
if self.messages:
|
|
return self.messages.pop(0)
|
|
return None
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
class FakeRedis:
|
|
def __init__(self):
|
|
self.sorted_sets = defaultdict(dict)
|
|
self.hashes = defaultdict(dict)
|
|
self.sets = defaultdict(set)
|
|
self.published = []
|
|
self.pubsub_instance = FakePubSub()
|
|
|
|
def pipeline(self):
|
|
return FakePipeline(self)
|
|
|
|
def zadd(self, key, mapping):
|
|
self.sorted_sets[key].update(mapping)
|
|
return len(mapping)
|
|
|
|
def hset(self, key, field, value):
|
|
self.hashes[key][field] = value
|
|
return 1
|
|
|
|
def sadd(self, key, *members):
|
|
before = len(self.sets[key])
|
|
self.sets[key].update(members)
|
|
return len(self.sets[key]) - before
|
|
|
|
def zrevrange(self, key, start, stop):
|
|
items = sorted(
|
|
self.sorted_sets[key].items(),
|
|
key=lambda item: (item[1], item[0]),
|
|
reverse=True,
|
|
)
|
|
if stop == -1:
|
|
return [member for member, _ in items[start:]]
|
|
return [member for member, _ in items[start : stop + 1]]
|
|
|
|
def hget(self, key, field):
|
|
return self.hashes[key].get(field)
|
|
|
|
def zrem(self, key, *members):
|
|
removed = 0
|
|
for member in members:
|
|
if member in self.sorted_sets[key]:
|
|
del self.sorted_sets[key][member]
|
|
removed += 1
|
|
return removed
|
|
|
|
def hdel(self, key, *fields):
|
|
removed = 0
|
|
for field in fields:
|
|
if field in self.hashes[key]:
|
|
del self.hashes[key][field]
|
|
removed += 1
|
|
return removed
|
|
|
|
def smembers(self, key):
|
|
return set(self.sets[key])
|
|
|
|
def srem(self, key, member):
|
|
if member in self.sets[key]:
|
|
self.sets[key].remove(member)
|
|
return 1
|
|
return 0
|
|
|
|
def zrangebyscore(self, key, min_score, max_score):
|
|
lower = float("-inf") if min_score == "-inf" else float(min_score)
|
|
upper = float(max_score)
|
|
return [
|
|
member
|
|
for member, score in self.sorted_sets[key].items()
|
|
if lower <= score <= upper
|
|
]
|
|
|
|
def zcard(self, key):
|
|
return len(self.sorted_sets[key])
|
|
|
|
def publish(self, channel, message):
|
|
self.published.append((channel, json.loads(message)))
|
|
return 1
|
|
|
|
def pubsub(self, ignore_subscribe_messages=True):
|
|
return self.pubsub_instance
|