Add TIMEOUT to query class (#2519)

* add timeout to query class

* Add test_timeout

* fix lines

* fix format

* add test & fixes

* merge tests

* change timeout to not_a_number

* change q1 to q2

* Fix async method

2 files changed, 18 insertions, 0 deletions

diff --git a/tests/test_asyncio/test_search.py b/tests/test_asyncio/test_search.py
index 88c80d3..12313b6 100644
--- a/tests/test_asyncio/test_search.py
+++ b/tests/test_asyncio/test_search.py
@@ -1001,3 +1001,12 @@ async def test_search_commands_in_pipeline(modclient: redis.Redis):
     assert "doc2" == res[3][4]
     assert res[3][5] is None
     assert res[3][3] == res[3][6] == ["txt", "foo bar"]
+
+
+@pytest.mark.redismod
+async def test_query_timeout(modclient: redis.Redis):
+    q1 = Query("foo").timeout(5000)
+    assert q1.get_args() == ["foo", "TIMEOUT", 5000, "LIMIT", 0, 10]
+    q2 = Query("foo").timeout("not_a_number")
+    with pytest.raises(redis.ResponseError):
+        await modclient.ft().search(q2)
diff --git a/tests/test_search.py b/tests/test_search.py
index 08c76f7..12876f6 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -1615,3 +1615,12 @@ def test_withsuffixtrie(modclient: redis.Redis):
     waitForIndex(modclient, getattr(modclient.ft(), "index_name", "idx"))
     info = modclient.ft().info()
     assert "WITHSUFFIXTRIE" in info["attributes"][0]
+
+
+@pytest.mark.redismod
+def test_query_timeout(modclient: redis.Redis):
+    q1 = Query("foo").timeout(5000)
+    assert q1.get_args() == ["foo", "TIMEOUT", 5000, "LIMIT", 0, 10]
+    q2 = Query("foo").timeout("not_a_number")
+    with pytest.raises(redis.ResponseError):
+        modclient.ft().search(q2)