|
9 | 9 | import re |
10 | 10 | import sys |
11 | 11 | from string import ascii_letters |
| 12 | +from typing import Any, Dict, List |
12 | 13 |
|
13 | 14 | import pytest |
14 | 15 | import pytest_asyncio |
@@ -2484,6 +2485,248 @@ async def test_geopos_no_value(self, r: valkey.Valkey): |
2484 | 2485 | async def test_old_geopos_no_value(self, r: valkey.Valkey): |
2485 | 2486 | assert await r.geopos("barcelona", "place1", "place2") == [] |
2486 | 2487 |
|
| 2488 | + @skip_if_server_version_lt("6.2.0") |
| 2489 | + async def test_geosearch(self, r: valkey.Valkey): |
| 2490 | + values = ( |
| 2491 | + (2.1909389952632, 41.433791470673, "place1") |
| 2492 | + + (2.1873744593677, 41.406342043777, b"\x80place2") |
| 2493 | + + (2.583333, 41.316667, "place3") |
| 2494 | + ) |
| 2495 | + await r.geoadd("barcelona", values) |
| 2496 | + assert await r.geosearch( |
| 2497 | + "barcelona", longitude=2.191, latitude=41.433, radius=1000 |
| 2498 | + ) == [b"place1"] |
| 2499 | + assert await r.geosearch( |
| 2500 | + "barcelona", longitude=2.187, latitude=41.406, radius=1000 |
| 2501 | + ) == [b"\x80place2"] |
| 2502 | + assert await r.geosearch( |
| 2503 | + "barcelona", longitude=2.191, latitude=41.433, height=1000, width=1000 |
| 2504 | + ) == [b"place1"] |
| 2505 | + assert await r.geosearch( |
| 2506 | + "barcelona", member="place3", radius=100, unit="km" |
| 2507 | + ) == [ |
| 2508 | + b"\x80place2", |
| 2509 | + b"place1", |
| 2510 | + b"place3", |
| 2511 | + ] |
| 2512 | + # test count |
| 2513 | + assert await r.geosearch( |
| 2514 | + "barcelona", member="place3", radius=100, unit="km", count=2 |
| 2515 | + ) == [b"place3", b"\x80place2"] |
| 2516 | + search_res = await r.geosearch( |
| 2517 | + "barcelona", member="place3", radius=100, unit="km", count=1, any=1 |
| 2518 | + ) |
| 2519 | + assert search_res[0] in [b"place1", b"place3", b"\x80place2"] |
| 2520 | + |
| 2521 | + @skip_unless_arch_bits(64) |
| 2522 | + @skip_if_server_version_lt("6.2.0") |
| 2523 | + async def test_geosearch_member(self, r: valkey.Valkey): |
| 2524 | + values = (2.1909389952632, 41.433791470673, "place1") + ( |
| 2525 | + 2.1873744593677, |
| 2526 | + 41.406342043777, |
| 2527 | + b"\x80place2", |
| 2528 | + ) |
| 2529 | + |
| 2530 | + await r.geoadd("barcelona", values) |
| 2531 | + assert await r.geosearch("barcelona", member="place1", radius=10) == [b"place1"] |
| 2532 | + |
| 2533 | + geosearch_place2, geosearch_place1 = await r.geosearch( |
| 2534 | + "barcelona", |
| 2535 | + member="place1", |
| 2536 | + radius=4000, |
| 2537 | + withdist=True, |
| 2538 | + withcoord=True, |
| 2539 | + withhash=True, |
| 2540 | + ) |
| 2541 | + |
| 2542 | + # All but the coordinates are identical |
| 2543 | + assert geosearch_place2[:-1] == [ |
| 2544 | + b"\x80place2", |
| 2545 | + 3067.4157, |
| 2546 | + 3471609625421029, |
| 2547 | + ] |
| 2548 | + assert_geo_is_close( |
| 2549 | + geosearch_place2[-1], (2.187376320362091, 41.40634178640635) |
| 2550 | + ) |
| 2551 | + assert geosearch_place1[:-1] == [ |
| 2552 | + b"place1", |
| 2553 | + 0.0, |
| 2554 | + 3471609698139488, |
| 2555 | + ] |
| 2556 | + assert_geo_is_close( |
| 2557 | + geosearch_place1[-1], (2.1909382939338684, 41.433790281840835) |
| 2558 | + ) |
| 2559 | + |
| 2560 | + @skip_if_server_version_lt("6.2.0") |
| 2561 | + async def test_geosearch_sort(self, r: valkey.Valkey): |
| 2562 | + values = (2.1909389952632, 41.433791470673, "place1") + ( |
| 2563 | + 2.1873744593677, |
| 2564 | + 41.406342043777, |
| 2565 | + "place2", |
| 2566 | + ) |
| 2567 | + await r.geoadd("barcelona", values) |
| 2568 | + assert await r.geosearch( |
| 2569 | + "barcelona", longitude=2.191, latitude=41.433, radius=3000, sort="ASC" |
| 2570 | + ) == [b"place1", b"place2"] |
| 2571 | + assert await r.geosearch( |
| 2572 | + "barcelona", longitude=2.191, latitude=41.433, radius=3000, sort="DESC" |
| 2573 | + ) == [b"place2", b"place1"] |
| 2574 | + |
| 2575 | + @skip_unless_arch_bits(64) |
| 2576 | + @skip_if_server_version_lt("6.2.0") |
| 2577 | + @pytest.mark.parametrize( |
| 2578 | + "geosearch_kwargs, expected_geosearch_result", |
| 2579 | + [ |
| 2580 | + ( |
| 2581 | + {"withdist": True, "withcoord": True, "withhash": True}, |
| 2582 | + [b"place1", 0.0881, 3471609698139488], |
| 2583 | + ), |
| 2584 | + ( |
| 2585 | + {"withdist": True, "withcoord": True}, |
| 2586 | + [b"place1", 0.0881], |
| 2587 | + ), |
| 2588 | + ( |
| 2589 | + {"withhash": True, "withcoord": True}, |
| 2590 | + [b"place1", 3471609698139488], |
| 2591 | + ), |
| 2592 | + ( |
| 2593 | + {"withdist": True, "withhash": True}, |
| 2594 | + [b"place1", 0.0881, 3471609698139488], |
| 2595 | + ), |
| 2596 | + ], |
| 2597 | + ) |
| 2598 | + async def test_geosearch_with( |
| 2599 | + self, |
| 2600 | + r: valkey.Valkey, |
| 2601 | + geosearch_kwargs: Dict[str, Any], |
| 2602 | + expected_geosearch_result: List[Any], |
| 2603 | + ): |
| 2604 | + values = (2.1909389952632, 41.433791470673, "place1") + ( |
| 2605 | + 2.1873744593677, |
| 2606 | + 41.406342043777, |
| 2607 | + "place2", |
| 2608 | + ) |
| 2609 | + await r.geoadd("barcelona", values) |
| 2610 | + |
| 2611 | + # test a bunch of combinations to test the parse response |
| 2612 | + # function. |
| 2613 | + geosearch_result = await r.geosearch( |
| 2614 | + "barcelona", |
| 2615 | + longitude=2.191, |
| 2616 | + latitude=41.433, |
| 2617 | + radius=1, |
| 2618 | + unit="km", |
| 2619 | + **geosearch_kwargs, |
| 2620 | + ) |
| 2621 | + assert len(geosearch_result) == 1 |
| 2622 | + if "withcoord" in geosearch_kwargs: |
| 2623 | + assert_geo_is_close( |
| 2624 | + geosearch_result[0][-1], (2.1909382939338684, 41.433790281840835) |
| 2625 | + ) |
| 2626 | + assert geosearch_result[0][:-1] == expected_geosearch_result |
| 2627 | + else: |
| 2628 | + assert geosearch_result == [expected_geosearch_result] |
| 2629 | + |
| 2630 | + assert ( |
| 2631 | + await r.geosearch( |
| 2632 | + "barcelona", |
| 2633 | + longitude=2, |
| 2634 | + latitude=1, |
| 2635 | + radius=1, |
| 2636 | + unit="km", |
| 2637 | + **geosearch_kwargs, |
| 2638 | + ) |
| 2639 | + == [] |
| 2640 | + ) |
| 2641 | + |
| 2642 | + @skip_if_server_version_lt("6.2.0") |
| 2643 | + async def test_geosearch_negative(self, r: valkey.Valkey): |
| 2644 | + # not specifying member nor longitude and latitude |
| 2645 | + with pytest.raises(exceptions.DataError): |
| 2646 | + assert await r.geosearch("barcelona") |
| 2647 | + # specifying member and longitude and latitude |
| 2648 | + with pytest.raises(exceptions.DataError): |
| 2649 | + assert await r.geosearch( |
| 2650 | + "barcelona", member="Paris", longitude=2, latitude=1 |
| 2651 | + ) |
| 2652 | + # specifying one of longitude and latitude |
| 2653 | + with pytest.raises(exceptions.DataError): |
| 2654 | + assert await r.geosearch("barcelona", longitude=2) |
| 2655 | + with pytest.raises(exceptions.DataError): |
| 2656 | + assert await r.geosearch("barcelona", latitude=2) |
| 2657 | + |
| 2658 | + # not specifying radius nor width and height |
| 2659 | + with pytest.raises(exceptions.DataError): |
| 2660 | + assert await r.geosearch("barcelona", member="Paris") |
| 2661 | + # specifying radius and width and height |
| 2662 | + with pytest.raises(exceptions.DataError): |
| 2663 | + assert await r.geosearch( |
| 2664 | + "barcelona", member="Paris", radius=3, width=2, height=1 |
| 2665 | + ) |
| 2666 | + # specifying one of width and height |
| 2667 | + with pytest.raises(exceptions.DataError): |
| 2668 | + assert await r.geosearch("barcelona", member="Paris", width=2) |
| 2669 | + with pytest.raises(exceptions.DataError): |
| 2670 | + assert await r.geosearch("barcelona", member="Paris", height=2) |
| 2671 | + |
| 2672 | + # invalid sort |
| 2673 | + with pytest.raises(exceptions.DataError): |
| 2674 | + assert await r.geosearch( |
| 2675 | + "barcelona", member="Paris", width=2, height=2, sort="wrong" |
| 2676 | + ) |
| 2677 | + |
| 2678 | + # invalid unit |
| 2679 | + with pytest.raises(exceptions.DataError): |
| 2680 | + assert await r.geosearch( |
| 2681 | + "barcelona", member="Paris", width=2, height=2, unit="miles" |
| 2682 | + ) |
| 2683 | + |
| 2684 | + # use any without count |
| 2685 | + with pytest.raises(exceptions.DataError): |
| 2686 | + assert await r.geosearch("barcelona", member="place3", radius=100, any=1) |
| 2687 | + |
| 2688 | + @pytest.mark.onlynoncluster |
| 2689 | + @skip_if_server_version_lt("6.2.0") |
| 2690 | + async def test_geosearchstore(self, r: valkey.Valkey): |
| 2691 | + values = (2.1909389952632, 41.433791470673, "place1") + ( |
| 2692 | + 2.1873744593677, |
| 2693 | + 41.406342043777, |
| 2694 | + "place2", |
| 2695 | + ) |
| 2696 | + |
| 2697 | + await r.geoadd("barcelona", values) |
| 2698 | + await r.geosearchstore( |
| 2699 | + "places_barcelona", |
| 2700 | + "barcelona", |
| 2701 | + longitude=2.191, |
| 2702 | + latitude=41.433, |
| 2703 | + radius=1000, |
| 2704 | + ) |
| 2705 | + assert await r.zrange("places_barcelona", 0, -1) == [b"place1"] |
| 2706 | + |
| 2707 | + @pytest.mark.onlynoncluster |
| 2708 | + @skip_unless_arch_bits(64) |
| 2709 | + @skip_if_server_version_lt("6.2.0") |
| 2710 | + async def test_geosearchstore_dist(self, r: valkey.Valkey): |
| 2711 | + values = (2.1909389952632, 41.433791470673, "place1") + ( |
| 2712 | + 2.1873744593677, |
| 2713 | + 41.406342043777, |
| 2714 | + "place2", |
| 2715 | + ) |
| 2716 | + |
| 2717 | + await r.geoadd("barcelona", values) |
| 2718 | + await r.geosearchstore( |
| 2719 | + "places_barcelona", |
| 2720 | + "barcelona", |
| 2721 | + longitude=2.191, |
| 2722 | + latitude=41.433, |
| 2723 | + radius=1000, |
| 2724 | + storedist=True, |
| 2725 | + ) |
| 2726 | + # instead of save the geo score, the distance is saved. |
| 2727 | + score = await r.zscore("places_barcelona", "place1") |
| 2728 | + assert math.isclose(score, 88.05060698409301) |
| 2729 | + |
2487 | 2730 | @skip_if_server_version_lt("3.2.0") |
2488 | 2731 | async def test_georadius(self, r: valkey.Valkey): |
2489 | 2732 | values = (2.1909389952632, 41.433791470673, "place1") + ( |
|
0 commit comments