File: test_zset.py

package info (click to toggle)
python-fakeredis 2.29.0-4
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 1,772 kB
  • sloc: python: 19,002; sh: 8; makefile: 5
file content (85 lines) | stat: -rw-r--r-- 3,053 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
from .base import (
    BaseTest,
    scores,
    fields,
    commands,
    st,
    keys,
    common_commands,
    counts,
    zero_or_more,
    optional,
    string_tests,
    float_as_bytes,
    Command,
)

limits = st.just(()) | st.tuples(st.just("limit"), counts, counts)
score_tests = scores | st.builds(lambda x: b"(" + repr(x).encode(), scores)
zset_no_score_create_commands = commands(st.just("zadd"), keys, st.lists(st.tuples(st.just(0), fields), min_size=1))
zset_no_score_commands = (  # TODO: test incr
    commands(
        st.just("zadd"),
        keys,
        *zero_or_more("nx", "xx", "ch", "incr"),
        st.lists(st.tuples(st.just(0), fields)),
    )
    | commands(st.just("zlexcount"), keys, string_tests, string_tests)
    | commands(st.sampled_from(["zrangebylex", "zrevrangebylex"]), keys, string_tests, string_tests, limits)
    | commands(st.just("zremrangebylex"), keys, string_tests, string_tests)
)


def build_zstore(command, dest, sources, weights, aggregate) -> Command:
    args = [command, dest, len(sources)]
    args += [source[0] for source in sources]
    if weights:
        args.append("weights")
        args += [source[1] for source in sources]
    if aggregate:
        args += ["aggregate", aggregate]
    return Command(args)


class TestZSet(BaseTest):
    zset_commands = (
        commands(
            st.just("zadd"),
            keys,
            *zero_or_more("nx", "xx", "ch", "incr"),
            st.lists(st.tuples(scores, fields)),
        )
        | commands(st.just("zcard"), keys)
        | commands(st.just("zcount"), keys, score_tests, score_tests)
        | commands(st.just("zincrby"), keys, scores, fields)
        | commands(st.sampled_from(["zrange", "zrevrange"]), keys, counts, counts, optional("withscores"))
        | commands(
            st.sampled_from(["zrangebyscore", "zrevrangebyscore"]),
            keys,
            score_tests,
            score_tests,
            limits,
            optional("withscores"),
        )
        | commands(st.sampled_from(["zrank", "zrevrank"]), keys, fields)
        | commands(st.just("zrem"), keys, st.lists(fields))
        | commands(st.just("zremrangebyrank"), keys, counts, counts)
        | commands(st.just("zremrangebyscore"), keys, score_tests, score_tests)
        | commands(st.just("zscore"), keys, fields)
        | st.builds(
            build_zstore,
            command=st.sampled_from(["zunionstore", "zinterstore"]),
            dest=keys,
            sources=st.lists(st.tuples(keys, float_as_bytes)),
            weights=st.booleans(),
            aggregate=st.sampled_from([None, "sum", "min", "max"]),
        )
    )
    # TODO: zscan, zpopmin/zpopmax, bzpopmin/bzpopmax, probably more
    create_command_strategy = commands(st.just("zadd"), keys, st.lists(st.tuples(scores, fields), min_size=1))
    command_strategy = zset_commands | common_commands


class TestZSetNoScores(BaseTest):
    create_command_strategy = zset_no_score_create_commands
    command_strategy = zset_no_score_commands | common_commands