diff --git a/tests/test_search.py b/tests/test_search.py index 6b40c5129d..7d335f6cec 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -2105,16 +2105,60 @@ def test_geo_params(client): params_dict = {"lat": "34.95126", "lon": "29.69465", "radius": 1000, "units": "km"} q = Query("@g:[$lon $lat $radius $units]").dialect(2) res = client.ft().search(q, query_params=params_dict) - if is_resp2_connection(client): - assert 3 == res.total - assert "doc1" == res.docs[0].id - assert "doc2" == res.docs[1].id - assert "doc3" == res.docs[2].id - else: - assert 3 == res["total_results"] - assert "doc1" == res["results"][0]["id"] - assert "doc2" == res["results"][1]["id"] - assert "doc3" == res["results"][2]["id"] + _assert_geosearch_result(client, res, ["doc1", "doc2", "doc3"]) + + +@pytest.mark.redismod +def test_geoshapes_query_intersects_and_disjoint(client): + client.ft().create_index((GeoShapeField("g", coord_system=GeoShapeField.FLAT))) + client.hset("doc_point1", mapping={"g": "POINT (10 10)"}) + client.hset("doc_point2", mapping={"g": "POINT (50 50)"}) + client.hset("doc_polygon1", mapping={"g": "POLYGON ((20 20, 25 35, 35 25, 20 20))"}) + client.hset( + "doc_polygon2", mapping={"g": "POLYGON ((60 60, 65 75, 70 70, 65 55, 60 60))"} + ) + + intersection = client.ft().search( + Query("@g:[intersects $shape]").dialect(3), + query_params={"shape": "POLYGON((15 15, 75 15, 50 70, 20 40, 15 15))"}, + ) + _assert_geosearch_result(client, intersection, ["doc_point2", "doc_polygon1"]) + + disjunction = client.ft().search( + Query("@g:[disjoint $shape]").dialect(3), + query_params={"shape": "POLYGON((15 15, 75 15, 50 70, 20 40, 15 15))"}, + ) + _assert_geosearch_result(client, disjunction, ["doc_point1", "doc_polygon2"]) + + +@pytest.mark.redismod +@skip_ifmodversion_lt("2.10.0", "search") +def test_geoshapes_query_contains_and_within(client): + client.ft().create_index((GeoShapeField("g", coord_system=GeoShapeField.FLAT))) + client.hset("doc_point1", mapping={"g": "POINT (10 10)"}) + client.hset("doc_point2", mapping={"g": "POINT (50 50)"}) + client.hset("doc_polygon1", mapping={"g": "POLYGON ((20 20, 25 35, 35 25, 20 20))"}) + client.hset( + "doc_polygon2", mapping={"g": "POLYGON ((60 60, 65 75, 70 70, 65 55, 60 60))"} + ) + + contains_a = client.ft().search( + Query("@g:[contains $shape]").dialect(3), + query_params={"shape": "POINT(25 25)"}, + ) + _assert_geosearch_result(client, contains_a, ["doc_polygon1"]) + + contains_b = client.ft().search( + Query("@g:[contains $shape]").dialect(3), + query_params={"shape": "POLYGON((24 24, 24 26, 25 25, 24 24))"}, + ) + _assert_geosearch_result(client, contains_b, ["doc_polygon1"]) + + within = client.ft().search( + Query("@g:[within $shape]").dialect(3), + query_params={"shape": "POLYGON((15 15, 75 15, 50 70, 20 40, 15 15))"}, + ) + _assert_geosearch_result(client, within, ["doc_point2", "doc_polygon1"]) @pytest.mark.redismod @@ -2278,7 +2322,19 @@ def test_geoshape(client: redis.Redis): q2 = Query("@geom:[CONTAINS $poly]").dialect(3) qp2 = {"poly": "POLYGON((2 2, 2 50, 50 50, 50 2, 2 2))"} result = client.ft().search(q1, query_params=qp1) - assert len(result.docs) == 1 - assert result.docs[0]["id"] == "small" + _assert_geosearch_result(client, result, ["small"]) result = client.ft().search(q2, query_params=qp2) - assert len(result.docs) == 2 + _assert_geosearch_result(client, result, ["small", "large"]) + + +def _assert_geosearch_result(client, result, expected_doc_ids): + """ + Make sure the result of a geo search is as expected, taking into account the RESP + version being used. + """ + if is_resp2_connection(client): + assert set([doc.id for doc in result.docs]) == set(expected_doc_ids) + assert result.total == len(expected_doc_ids) + else: + assert set([doc["id"] for doc in result["results"]]) == set(expected_doc_ids) + assert result["total_results"] == len(expected_doc_ids)