async def test_negate_auto_close_client_pool(self, r: redis.Redis,
auto_close_conn_pool):
r.auto_close_connection_pool = auto_close_conn_pool
new_conn = await self.create_two_conn(r)
await r.close(close_connection_pool=False)
assert not self.has_no_connected_connections(r.connection_pool)
assert r.connection_pool._in_use_connections == {new_conn}
assert r.connection_pool._available_connections[0].is_connected
assert self.get_total_connected_connections(r.connection_pool) == 2
async def test_connection_error_raised_when_connection_dies(self, r: redis.Redis):
p = r.pubsub()
await p.subscribe("foo")
assert await wait_for_message(p) == make_message("subscribe", "foo", 1)
for client in await r.client_list():
if client["cmd"] == "subscribe":
await r.client_kill_filter(_id=client["id"])
with pytest.raises(ConnectionError):
await wait_for_message(p)
async def test_unicode_channel_message_handler(self, r: redis.Redis):
p = r.pubsub(ignore_subscribe_messages=True)
channel = "uni" + chr(4456) + "code"
channels = {channel: self.message_handler}
await p.subscribe(**channels)
assert await wait_for_message(p) is None
assert await r.publish(channel, "test message") == 1
assert await wait_for_message(p) is None
assert self.message == make_message("message", channel, "test message")
async def test_pattern_message_handler(self, r: redis.Redis):
p = r.pubsub(ignore_subscribe_messages=True)
await p.psubscribe(**{"f*": self.message_handler})
assert await wait_for_message(p) is None
assert await r.publish("foo", "test message") == 1
assert await wait_for_message(p) is None
assert self.message == make_message(
"pmessage", "foo", "test message", pattern="f*"
)
async def test_published_message_to_channel(self, r: redis.Redis):
p = r.pubsub()
await p.subscribe("foo")
assert await wait_for_message(p) == make_message("subscribe", "foo", 1)
assert await r.publish("foo", "test message") == 1
message = await wait_for_message(p)
assert isinstance(message, dict)
assert message == make_message("message", "foo", "test message")
async def test_pattern_publish(self, r: redis.Redis):
p = r.pubsub()
await p.psubscribe(self.pattern)
assert await wait_for_message(p) == self.make_message(
"psubscribe", self.pattern, 1
)
await r.publish(self.channel, self.data)
assert await wait_for_message(p) == self.make_message(
"pmessage", self.channel, self.data, pattern=self.pattern
)
async def test_channel_publish(self, r: redis.Redis):
p = r.pubsub()
await p.subscribe(self.channel)
assert await wait_for_message(p) == self.make_message(
"subscribe", self.channel, 1
)
await r.publish(self.channel, self.data)
assert await wait_for_message(p) == self.make_message(
"message", self.channel, self.data
)
async def test_pattern_subscribe_unsubscribe(self, r: redis.Redis):
p = r.pubsub()
await p.psubscribe(self.pattern)
assert await wait_for_message(p) == self.make_message(
"psubscribe", self.pattern, 1
)
await p.punsubscribe(self.pattern)
assert await wait_for_message(p) == self.make_message(
"punsubscribe", self.pattern, 0
)
async def test_channel_subscribe_unsubscribe(self, r: redis.Redis):
p = r.pubsub()
await p.subscribe(self.channel)
assert await wait_for_message(p) == self.make_message(
"subscribe", self.channel, 1
)
await p.unsubscribe(self.channel)
assert await wait_for_message(p) == self.make_message(
"unsubscribe", self.channel, 0
)
async def test_unicode_pattern_message_handler(self, r: redis.Redis):
p = r.pubsub(ignore_subscribe_messages=True)
pattern = "uni" + chr(4456) + "*"
channel = "uni" + chr(4456) + "code"
await p.psubscribe(**{pattern: self.message_handler})
assert await wait_for_message(p) is None
assert await r.publish(channel, "test message") == 1
assert await wait_for_message(p) is None
assert self.message == make_message(
"pmessage", channel, "test message", pattern=pattern
)
async def test_filters(modclient: redis.Redis):
await (
modclient.ft().create_index(
(TextField("txt"), NumericField("num"), GeoField("loc"))
)
)
await (
modclient.ft().add_document(
"doc1", txt="foo bar", num=3.141, loc="-0.441,51.458"
)
)
await modclient.ft().add_document("doc2", txt="foo baz", num=2, loc="-0.1,51.2")
await waitForIndex(modclient, "idx")
# Test numerical filter
q1 = Query("foo").add_filter(NumericFilter("num", 0, 2)).no_content()
q2 = (
Query("foo")
.add_filter(NumericFilter("num", 2, NumericFilter.INF, minExclusive=True))
.no_content()
)
res1, res2 = await modclient.ft().search(q1), await modclient.ft().search(q2)
assert 1 == res1.total
assert 1 == res2.total
assert "doc2" == res1.docs[0].id
assert "doc1" == res2.docs[0].id
# Test geo filter
q1 = Query("foo").add_filter(GeoFilter("loc", -0.44, 51.45, 10)).no_content()
q2 = Query("foo").add_filter(GeoFilter("loc", -0.44, 51.45, 100)).no_content()
res1, res2 = await modclient.ft().search(q1), await modclient.ft().search(q2)
assert 1 == res1.total
assert 2 == res2.total
assert "doc1" == res1.docs[0].id
# Sort results, after RDB reload order may change
res = [res2.docs[0].id, res2.docs[1].id]
res.sort()
assert ["doc1", "doc2"] == res
async def test_no_create(modclient: redis.Redis):
await (
modclient.ft().create_index((TextField("f1"), TextField("f2"), TextField("f3")))
)
await modclient.ft().add_document("doc1", f1="f1_val", f2="f2_val")
await modclient.ft().add_document("doc2", f1="f1_val", f2="f2_val")
await modclient.ft().add_document("doc1", f3="f3_val", no_create=True)
await modclient.ft().add_document("doc2", f3="f3_val", no_create=True, partial=True)
await waitForIndex(modclient, "idx")
# Search for f3 value. All documents should have it
res = await modclient.ft().search("@f3:f3_val")
assert 2 == res.total
# Only the document updated with PARTIAL should still have f1 and f2 values
res = await modclient.ft().search("@f3:f3_val @f2:f2_val @f1:f1_val")
assert 1 == res.total
with pytest.raises(redis.ResponseError):
await (
modclient.ft().add_document(
"doc3", f2="f2_val", f3="f3_val", no_create=True
)
)
async def test_auto_complete(modclient: redis.Redis):
n = 0
with open(TITLES_CSV) as f:
cr = csv.reader(f)
for row in cr:
n += 1
term, score = row[0], float(row[1])
assert n == await modclient.ft().sugadd("ac", Suggestion(term, score=score))
assert n == await modclient.ft().suglen("ac")
ret = await modclient.ft().sugget("ac", "bad", with_scores=True)
assert 2 == len(ret)
assert "badger" == ret[0].string
assert isinstance(ret[0].score, float)
assert 1.0 != ret[0].score
assert "badalte rishtey" == ret[1].string
assert isinstance(ret[1].score, float)
assert 1.0 != ret[1].score
ret = await modclient.ft().sugget("ac", "bad", fuzzy=True, num=10)
assert 10 == len(ret)
assert 1.0 == ret[0].score
strs = {x.string for x in ret}
for sug in strs:
assert 1 == await modclient.ft().sugdel("ac", sug)
# make sure a second delete returns 0
for sug in strs:
assert 0 == await modclient.ft().sugdel("ac", sug)
# make sure they were actually deleted
ret2 = await modclient.ft().sugget("ac", "bad", fuzzy=True, num=10)
for sug in ret2:
assert sug.string not in strs
# Test with payload
await modclient.ft().sugadd("ac", Suggestion("pay1", payload="pl1"))
await modclient.ft().sugadd("ac", Suggestion("pay2", payload="pl2"))
await modclient.ft().sugadd("ac", Suggestion("pay3", payload="pl3"))
sugs = await (
modclient.ft().sugget("ac", "pay", with_payloads=True, with_scores=True)
)
assert 3 == len(sugs)
for sug in sugs:
assert sug.payload
assert sug.payload.startswith("pl")
async def test_ignore_all_subscribe_messages(self, r: redis.Redis):
p = r.pubsub(ignore_subscribe_messages=True)
checks = (
(p.subscribe, "foo"),
(p.unsubscribe, "foo"),
(p.psubscribe, "f*"),
(p.punsubscribe, "f*"),
)
assert p.subscribed is False
for func, channel in checks:
assert await func(channel) is None
assert p.subscribed is True
assert await wait_for_message(p) is None
assert p.subscribed is False
async def test_channel_message_handler(self, r: redis.Redis):
p = r.pubsub(ignore_subscribe_messages=True)
await p.subscribe(**{self.channel: self.message_handler})
assert await wait_for_message(p) is None
await r.publish(self.channel, self.data)
assert await wait_for_message(p) is None
assert self.message == self.make_message("message", self.channel, self.data)
# test that we reconnected to the correct channel
self.message = None
await p.connection.disconnect()
assert await wait_for_message(p) is None # should reconnect
new_data = self.data + "new data"
await r.publish(self.channel, new_data)
assert await wait_for_message(p) is None
assert self.message == self.make_message("message", self.channel, new_data)
async def log_redis_info(redis: Redis, log_func: Callable[[str], Any]) -> None:
async with redis.pipeline(transaction=True) as pipe:
pipe.info(section='Server')
pipe.info(section='Memory')
pipe.info(section='Clients')
pipe.dbsize()
info_server, info_memory, info_clients, key_count = await pipe.execute(
)
redis_version = info_server.get('redis_version', '?')
mem_usage = info_memory.get('used_memory_human', '?')
clients_connected = info_clients.get('connected_clients', '?')
log_func(f'redis_version={redis_version} '
f'mem_usage={mem_usage} '
f'clients_connected={clients_connected} '
f'db_keys={key_count}')
async def test_example(modclient: redis.Redis):
# Creating the index definition and schema
await (
modclient.ft().create_index((TextField("title", weight=5.0), TextField("body")))
)
# Indexing a document
await modclient.ft().add_document(
"doc1",
title="RediSearch",
body="Redisearch impements a search engine on top of redis",
)
# Searching with complex parameters:
q = Query("search engine").verbatim().no_content().paging(0, 5)
res = await modclient.ft().search(q)
assert res is not None
async def test_exception_handler(self, r: redis.Redis):
def exception_handler_callback(e, pubsub) -> None:
assert pubsub == p
exceptions.put_nowait(e)
exceptions = asyncio.Queue()
p = r.pubsub()
await self._subscribe(p, foo=lambda x: None)
with mock.patch.object(p, "get_message", side_effect=Exception("error")):
task = asyncio.get_event_loop().create_task(
p.run(exception_handler=exception_handler_callback)
)
e = await exceptions.get()
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
assert str(e) == "error"
async def test_replace(modclient: redis.Redis):
await modclient.ft().create_index((TextField("txt"),))
await modclient.ft().add_document("doc1", txt="foo bar")
await modclient.ft().add_document("doc2", txt="foo bar")
await waitForIndex(modclient, "idx")
res = await modclient.ft().search("foo bar")
assert 2 == res.total
await (
modclient.ft().add_document("doc1", replace=True, txt="this is a replaced doc")
)
res = await modclient.ft().search("foo bar")
assert 1 == res.total
assert "doc2" == res.docs[0].id
res = await modclient.ft().search("replaced doc")
assert 1 == res.total
assert "doc1" == res.docs[0].id
async def test_callbacks(self, r: redis.Redis):
def callback(message):
messages.put_nowait(message)
messages = asyncio.Queue()
p = r.pubsub()
await self._subscribe(p, foo=callback)
task = asyncio.get_event_loop().create_task(p.run())
await r.publish("foo", "bar")
message = await messages.get()
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
assert message == {
"channel": b"foo",
"data": b"bar",
"pattern": None,
"type": "message",
}
async def test_published_message_to_pattern(self, r: redis.Redis):
p = r.pubsub()
await p.subscribe("foo")
await p.psubscribe("f*")
assert await wait_for_message(p) == make_message("subscribe", "foo", 1)
assert await wait_for_message(p) == make_message("psubscribe", "f*", 2)
# 1 to pattern, 1 to channel
assert await r.publish("foo", "test message") == 2
message1 = await wait_for_message(p)
message2 = await wait_for_message(p)
assert isinstance(message1, dict)
assert isinstance(message2, dict)
expected = [
make_message("message", "foo", "test message"),
make_message("pmessage", "foo", "test message", pattern="f*"),
]
assert message1 in expected
assert message2 in expected
assert message1 != message2
async def wait_for_command(client: redis.Redis,
monitor: Monitor,
command: str,
key: Union[str, None] = None):
# issue a command with a key name that's local to this process.
# if we find a command with our key before the command we're waiting
# for, something went wrong
if key is None:
# generate key
redis_version = REDIS_INFO["version"]
if Version(redis_version) >= Version("5.0.0"):
id_str = str(client.client_id())
else:
id_str = f"{random.randrange(2 ** 32):08x}"
key = f"__REDIS-PY-{id_str}__"
await client.get(key)
while True:
monitor_response = await monitor.next_command()
if command in monitor_response["command"]:
return monitor_response
if key in monitor_response["command"]:
return None
async def test_spell_check(modclient: redis.Redis):
await modclient.ft().create_index((TextField("f1"), TextField("f2")))
await (
modclient.ft().add_document(
"doc1", f1="some valid content", f2="this is sample text"
)
)
await modclient.ft().add_document("doc2", f1="very important", f2="lorem ipsum")
await waitForIndex(modclient, "idx")
# test spellcheck
res = await modclient.ft().spellcheck("impornant")
assert "important" == res["impornant"][0]["suggestion"]
res = await modclient.ft().spellcheck("contnt")
assert "content" == res["contnt"][0]["suggestion"]
# test spellcheck with Levenshtein distance
res = await modclient.ft().spellcheck("vlis")
assert res == {}
res = await modclient.ft().spellcheck("vlis", distance=2)
assert "valid" == res["vlis"][0]["suggestion"]
# test spellcheck include
await modclient.ft().dict_add("dict", "lore", "lorem", "lorm")
res = await modclient.ft().spellcheck("lorm", include="dict")
assert len(res["lorm"]) == 3
assert (
res["lorm"][0]["suggestion"],
res["lorm"][1]["suggestion"],
res["lorm"][2]["suggestion"],
) == ("lorem", "lore", "lorm")
assert (res["lorm"][0]["score"], res["lorm"][1]["score"]) == ("0.5", "0")
# test spellcheck exclude
res = await modclient.ft().spellcheck("lorm", exclude="dict")
assert res == {}
async def test_summarize(modclient: redis.Redis):
await createIndex(modclient.ft())
await waitForIndex(modclient, "idx")
q = Query("king henry").paging(0, 1)
q.highlight(fields=("play", "txt"), tags=("<b>", "</b>"))
q.summarize("txt")
doc = sorted((await modclient.ft().search(q)).docs)[0]
assert "<b>Henry</b> IV" == doc.play
assert (
"ACT I SCENE I. London. The palace. Enter <b>KING</b> <b>HENRY</b>, LORD JOHN OF LANCASTER, the EARL of WESTMORELAND, SIR... " # noqa
== doc.txt
)
q = Query("king henry").paging(0, 1).summarize().highlight()
doc = sorted((await modclient.ft().search(q)).docs)[0]
assert "<b>Henry</b> ... " == doc.play
assert (
"ACT I SCENE I. London. The palace. Enter <b>KING</b> <b>HENRY</b>, LORD JOHN OF LANCASTER, the EARL of WESTMORELAND, SIR... " # noqa
== doc.txt
)
async def test_sort_by(modclient: redis.Redis):
await (
modclient.ft().create_index(
(TextField("txt"), NumericField("num", sortable=True))
)
)
await modclient.ft().add_document("doc1", txt="foo bar", num=1)
await modclient.ft().add_document("doc2", txt="foo baz", num=2)
await modclient.ft().add_document("doc3", txt="foo qux", num=3)
# Test sort
q1 = Query("foo").sort_by("num", asc=True).no_content()
q2 = Query("foo").sort_by("num", asc=False).no_content()
res1, res2 = await modclient.ft().search(q1), await modclient.ft().search(q2)
assert 3 == res1.total
assert "doc1" == res1.docs[0].id
assert "doc2" == res1.docs[1].id
assert "doc3" == res1.docs[2].id
assert 3 == res2.total
assert "doc1" == res2.docs[2].id
assert "doc2" == res2.docs[1].id
assert "doc3" == res2.docs[0].id
async def test_pattern_subscribe_unsubscribe(self, r: redis.Redis):
kwargs = make_subscribe_test_data(r.pubsub(), "pattern")
await self._test_subscribe_unsubscribe(**kwargs)
async def test_channel_subscribe_unsubscribe(self, r: redis.Redis):
kwargs = make_subscribe_test_data(r.pubsub(), "channel")
await self._test_subscribe_unsubscribe(**kwargs)
async def test_get_message_with_timeout_returns_none(self, r: redis.Redis):
p = r.pubsub()
await p.subscribe("foo")
assert await wait_for_message(p) == make_message("subscribe", "foo", 1)
assert await p.get_message(timeout=0.01) is None
async def test_channel_subscribe(self, r: redis.Redis):
r = redis.Redis(host="localhost", port=6390)
p = r.pubsub()
with pytest.raises(ConnectionError):
await p.subscribe("foo")
async def test_pubsub_numpat(self, r: redis.Redis):
p = r.pubsub()
await p.psubscribe("*oo", "*ar", "b*z")
for i in range(3):
assert (await wait_for_message(p))["type"] == "psubscribe"
assert await r.pubsub_numpat() == 3
