path: root/tests/test_asyncio/test_bloom.py

from math import inf

import pytest

import redis.asyncio as redis
from redis.exceptions import ModuleError, RedisError
from redis.utils import HIREDIS_AVAILABLE
from tests.conftest import skip_ifmodversion_lt


def intlist(obj):
    return [int(v) for v in obj]


# @pytest.fixture
# async def client(modclient):
#     assert isinstance(modawait modclient.bf(), redis.commands.bf.BFBloom)
#     assert isinstance(modawait modclient.cf(), redis.commands.bf.CFBloom)
#     assert isinstance(modawait modclient.cms(), redis.commands.bf.CMSBloom)
#     assert isinstance(modawait modclient.tdigest(), redis.commands.bf.TDigestBloom)
#     assert isinstance(modawait modclient.topk(), redis.commands.bf.TOPKBloom)

#     modawait modclient.flushdb()
#     return modclient


@pytest.mark.redismod
async def test_create(modclient: redis.Redis):
    """Test CREATE/RESERVE calls"""
    assert await modclient.bf().create("bloom", 0.01, 1000)
    assert await modclient.bf().create("bloom_e", 0.01, 1000, expansion=1)
    assert await modclient.bf().create("bloom_ns", 0.01, 1000, noScale=True)
    assert await modclient.cf().create("cuckoo", 1000)
    assert await modclient.cf().create("cuckoo_e", 1000, expansion=1)
    assert await modclient.cf().create("cuckoo_bs", 1000, bucket_size=4)
    assert await modclient.cf().create("cuckoo_mi", 1000, max_iterations=10)
    assert await modclient.cms().initbydim("cmsDim", 100, 5)
    assert await modclient.cms().initbyprob("cmsProb", 0.01, 0.01)
    assert await modclient.topk().reserve("topk", 5, 100, 5, 0.9)


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_create(modclient: redis.Redis):
    assert await modclient.tdigest().create("tDigest", 100)


# region Test Bloom Filter
@pytest.mark.redismod
async def test_bf_add(modclient: redis.Redis):
    assert await modclient.bf().create("bloom", 0.01, 1000)
    assert 1 == await modclient.bf().add("bloom", "foo")
    assert 0 == await modclient.bf().add("bloom", "foo")
    assert [0] == intlist(await modclient.bf().madd("bloom", "foo"))
    assert [0, 1] == await modclient.bf().madd("bloom", "foo", "bar")
    assert [0, 0, 1] == await modclient.bf().madd("bloom", "foo", "bar", "baz")
    assert 1 == await modclient.bf().exists("bloom", "foo")
    assert 0 == await modclient.bf().exists("bloom", "noexist")
    assert [1, 0] == intlist(await modclient.bf().mexists("bloom", "foo", "noexist"))


@pytest.mark.redismod
async def test_bf_insert(modclient: redis.Redis):
    assert await modclient.bf().create("bloom", 0.01, 1000)
    assert [1] == intlist(await modclient.bf().insert("bloom", ["foo"]))
    assert [0, 1] == intlist(await modclient.bf().insert("bloom", ["foo", "bar"]))
    assert [1] == intlist(await modclient.bf().insert("captest", ["foo"], capacity=10))
    assert [1] == intlist(await modclient.bf().insert("errtest", ["foo"], error=0.01))
    assert 1 == await modclient.bf().exists("bloom", "foo")
    assert 0 == await modclient.bf().exists("bloom", "noexist")
    assert [1, 0] == intlist(await modclient.bf().mexists("bloom", "foo", "noexist"))
    info = await modclient.bf().info("bloom")
    assert 2 == info.insertedNum
    assert 1000 == info.capacity
    assert 1 == info.filterNum


@pytest.mark.redismod
async def test_bf_scandump_and_loadchunk(modclient: redis.Redis):
    # Store a filter
    await modclient.bf().create("myBloom", "0.0001", "1000")

    # test is probabilistic and might fail. It is OK to change variables if
    # certain to not break anything
    async def do_verify():
        res = 0
        for x in range(1000):
            await modclient.bf().add("myBloom", x)
            rv = await modclient.bf().exists("myBloom", x)
            assert rv
            rv = await modclient.bf().exists("myBloom", f"nonexist_{x}")
            res += rv == x
        assert res < 5

    await do_verify()
    cmds = []
    if HIREDIS_AVAILABLE:
        with pytest.raises(ModuleError):
            cur = await modclient.bf().scandump("myBloom", 0)
        return

    cur = await modclient.bf().scandump("myBloom", 0)
    first = cur[0]
    cmds.append(cur)

    while True:
        cur = await modclient.bf().scandump("myBloom", first)
        first = cur[0]
        if first == 0:
            break
        else:
            cmds.append(cur)
    prev_info = await modclient.bf().execute_command("bf.debug", "myBloom")

    # Remove the filter
    await modclient.bf().client.delete("myBloom")

    # Now, load all the commands:
    for cmd in cmds:
        await modclient.bf().loadchunk("myBloom", *cmd)

    cur_info = await modclient.bf().execute_command("bf.debug", "myBloom")
    assert prev_info == cur_info
    await do_verify()

    await modclient.bf().client.delete("myBloom")
    await modclient.bf().create("myBloom", "0.0001", "10000000")


@pytest.mark.redismod
async def test_bf_info(modclient: redis.Redis):
    expansion = 4
    # Store a filter
    await modclient.bf().create("nonscaling", "0.0001", "1000", noScale=True)
    info = await modclient.bf().info("nonscaling")
    assert info.expansionRate is None

    await modclient.bf().create("expanding", "0.0001", "1000", expansion=expansion)
    info = await modclient.bf().info("expanding")
    assert info.expansionRate == 4

    try:
        # noScale mean no expansion
        await modclient.bf().create(
            "myBloom", "0.0001", "1000", expansion=expansion, noScale=True
        )
        assert False
    except RedisError:
        assert True


@pytest.mark.redismod
async def test_bf_card(modclient: redis.Redis):
    # return 0 if the key does not exist
    assert await modclient.bf().card("not_exist") == 0

    # Store a filter
    assert await modclient.bf().add("bf1", "item_foo") == 1
    assert await modclient.bf().card("bf1") == 1

    # Error when key is of a type other than Bloom filter.
    with pytest.raises(redis.ResponseError):
        await modclient.set("setKey", "value")
        await modclient.bf().card("setKey")


# region Test Cuckoo Filter
@pytest.mark.redismod
async def test_cf_add_and_insert(modclient: redis.Redis):
    assert await modclient.cf().create("cuckoo", 1000)
    assert await modclient.cf().add("cuckoo", "filter")
    assert not await modclient.cf().addnx("cuckoo", "filter")
    assert 1 == await modclient.cf().addnx("cuckoo", "newItem")
    assert [1] == await modclient.cf().insert("captest", ["foo"])
    assert [1] == await modclient.cf().insert("captest", ["foo"], capacity=1000)
    assert [1] == await modclient.cf().insertnx("captest", ["bar"])
    assert [1] == await modclient.cf().insertnx("captest", ["food"], nocreate="1")
    assert [0, 0, 1] == await modclient.cf().insertnx("captest", ["foo", "bar", "baz"])
    assert [0] == await modclient.cf().insertnx("captest", ["bar"], capacity=1000)
    assert [1] == await modclient.cf().insert("empty1", ["foo"], capacity=1000)
    assert [1] == await modclient.cf().insertnx("empty2", ["bar"], capacity=1000)
    info = await modclient.cf().info("captest")
    assert 5 == info.insertedNum
    assert 0 == info.deletedNum
    assert 1 == info.filterNum


@pytest.mark.redismod
async def test_cf_exists_and_del(modclient: redis.Redis):
    assert await modclient.cf().create("cuckoo", 1000)
    assert await modclient.cf().add("cuckoo", "filter")
    assert await modclient.cf().exists("cuckoo", "filter")
    assert not await modclient.cf().exists("cuckoo", "notexist")
    assert 1 == await modclient.cf().count("cuckoo", "filter")
    assert 0 == await modclient.cf().count("cuckoo", "notexist")
    assert await modclient.cf().delete("cuckoo", "filter")
    assert 0 == await modclient.cf().count("cuckoo", "filter")


# region Test Count-Min Sketch
@pytest.mark.redismod
async def test_cms(modclient: redis.Redis):
    assert await modclient.cms().initbydim("dim", 1000, 5)
    assert await modclient.cms().initbyprob("prob", 0.01, 0.01)
    assert await modclient.cms().incrby("dim", ["foo"], [5])
    assert [0] == await modclient.cms().query("dim", "notexist")
    assert [5] == await modclient.cms().query("dim", "foo")
    assert [10, 15] == await modclient.cms().incrby("dim", ["foo", "bar"], [5, 15])
    assert [10, 15] == await modclient.cms().query("dim", "foo", "bar")
    info = await modclient.cms().info("dim")
    assert 1000 == info.width
    assert 5 == info.depth
    assert 25 == info.count


@pytest.mark.redismod
@pytest.mark.onlynoncluster
async def test_cms_merge(modclient: redis.Redis):
    assert await modclient.cms().initbydim("A", 1000, 5)
    assert await modclient.cms().initbydim("B", 1000, 5)
    assert await modclient.cms().initbydim("C", 1000, 5)
    assert await modclient.cms().incrby("A", ["foo", "bar", "baz"], [5, 3, 9])
    assert await modclient.cms().incrby("B", ["foo", "bar", "baz"], [2, 3, 1])
    assert [5, 3, 9] == await modclient.cms().query("A", "foo", "bar", "baz")
    assert [2, 3, 1] == await modclient.cms().query("B", "foo", "bar", "baz")
    assert await modclient.cms().merge("C", 2, ["A", "B"])
    assert [7, 6, 10] == await modclient.cms().query("C", "foo", "bar", "baz")
    assert await modclient.cms().merge("C", 2, ["A", "B"], ["1", "2"])
    assert [9, 9, 11] == await modclient.cms().query("C", "foo", "bar", "baz")
    assert await modclient.cms().merge("C", 2, ["A", "B"], ["2", "3"])
    assert [16, 15, 21] == await modclient.cms().query("C", "foo", "bar", "baz")


# endregion


# region Test Top-K
@pytest.mark.redismod
async def test_topk(modclient: redis.Redis):
    # test list with empty buckets
    assert await modclient.topk().reserve("topk", 3, 50, 4, 0.9)
    assert [
        None,
        None,
        None,
        "A",
        "C",
        "D",
        None,
        None,
        "E",
        None,
        "B",
        "C",
        None,
        None,
        None,
        "D",
        None,
    ] == await modclient.topk().add(
        "topk",
        "A",
        "B",
        "C",
        "D",
        "E",
        "A",
        "A",
        "B",
        "C",
        "G",
        "D",
        "B",
        "D",
        "A",
        "E",
        "E",
        1,
    )
    assert [1, 1, 0, 0, 1, 0, 0] == await modclient.topk().query(
        "topk", "A", "B", "C", "D", "E", "F", "G"
    )
    with pytest.deprecated_call():
        assert [4, 3, 2, 3, 3, 0, 1] == await modclient.topk().count(
            "topk", "A", "B", "C", "D", "E", "F", "G"
        )

    # test full list
    assert await modclient.topk().reserve("topklist", 3, 50, 3, 0.9)
    assert await modclient.topk().add(
        "topklist",
        "A",
        "B",
        "C",
        "D",
        "E",
        "A",
        "A",
        "B",
        "C",
        "G",
        "D",
        "B",
        "D",
        "A",
        "E",
        "E",
    )
    assert ["A", "B", "E"] == await modclient.topk().list("topklist")
    res = await modclient.topk().list("topklist", withcount=True)
    assert ["A", 4, "B", 3, "E", 3] == res
    info = await modclient.topk().info("topklist")
    assert 3 == info.k
    assert 50 == info.width
    assert 3 == info.depth
    assert 0.9 == round(float(info.decay), 1)


@pytest.mark.redismod
async def test_topk_incrby(modclient: redis.Redis):
    await modclient.flushdb()
    assert await modclient.topk().reserve("topk", 3, 10, 3, 1)
    assert [None, None, None] == await modclient.topk().incrby(
        "topk", ["bar", "baz", "42"], [3, 6, 2]
    )
    res = await modclient.topk().incrby("topk", ["42", "xyzzy"], [8, 4])
    assert [None, "bar"] == res
    with pytest.deprecated_call():
        assert [3, 6, 10, 4, 0] == await modclient.topk().count(
            "topk", "bar", "baz", "42", "xyzzy", 4
        )


# region Test T-Digest
@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_reset(modclient: redis.Redis):
    assert await modclient.tdigest().create("tDigest", 10)
    # reset on empty histogram
    assert await modclient.tdigest().reset("tDigest")
    # insert data-points into sketch
    assert await modclient.tdigest().add("tDigest", list(range(10)))

    assert await modclient.tdigest().reset("tDigest")
    # assert we have 0 unmerged nodes
    assert 0 == (await modclient.tdigest().info("tDigest")).unmerged_nodes


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_merge(modclient: redis.Redis):
    assert await modclient.tdigest().create("to-tDigest", 10)
    assert await modclient.tdigest().create("from-tDigest", 10)
    # insert data-points into sketch
    assert await modclient.tdigest().add("from-tDigest", [1.0] * 10)
    assert await modclient.tdigest().add("to-tDigest", [2.0] * 10)
    # merge from-tdigest into to-tdigest
    assert await modclient.tdigest().merge("to-tDigest", 1, "from-tDigest")
    # we should now have 110 weight on to-histogram
    info = await modclient.tdigest().info("to-tDigest")
    total_weight_to = float(info.merged_weight) + float(info.unmerged_weight)
    assert 20.0 == total_weight_to
    # test override
    assert await modclient.tdigest().create("from-override", 10)
    assert await modclient.tdigest().create("from-override-2", 10)
    assert await modclient.tdigest().add("from-override", [3.0] * 10)
    assert await modclient.tdigest().add("from-override-2", [4.0] * 10)
    assert await modclient.tdigest().merge(
        "to-tDigest", 2, "from-override", "from-override-2", override=True
    )
    assert 3.0 == await modclient.tdigest().min("to-tDigest")
    assert 4.0 == await modclient.tdigest().max("to-tDigest")


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_min_and_max(modclient: redis.Redis):
    assert await modclient.tdigest().create("tDigest", 100)
    # insert data-points into sketch
    assert await modclient.tdigest().add("tDigest", [1, 2, 3])
    # min/max
    assert 3 == await modclient.tdigest().max("tDigest")
    assert 1 == await modclient.tdigest().min("tDigest")


@pytest.mark.redismod
@pytest.mark.experimental
@skip_ifmodversion_lt("2.4.0", "bf")
async def test_tdigest_quantile(modclient: redis.Redis):
    assert await modclient.tdigest().create("tDigest", 500)
    # insert data-points into sketch
    assert await modclient.tdigest().add(
        "tDigest", list([x * 0.01 for x in range(1, 10000)])
    )
    # assert min min/max have same result as quantile 0 and 1
    assert (
        await modclient.tdigest().max("tDigest")
        == (await modclient.tdigest().quantile("tDigest", 1))[0]
    )
    assert (
        await modclient.tdigest().min("tDigest")
        == (await modclient.tdigest().quantile("tDigest", 0.0))[0]
    )

    assert 1.0 == round((await modclient.tdigest().quantile("tDigest", 0.01))[0], 2)
    assert 99.0 == round((await modclient.tdigest().quantile("tDigest", 0.99))[0], 2)

    # test multiple quantiles
    assert await modclient.tdigest().create("t-digest", 100)
    assert await modclient.tdigest().add("t-digest", [1, 2, 3, 4, 5])
    res = await modclient.tdigest().quantile("t-digest", 0.5, 0.8)
    assert [3.0, 5.0] == res


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_cdf(modclient: redis.Redis):
    assert await modclient.tdigest().create("tDigest", 100)
    # insert data-points into sketch
    assert await modclient.tdigest().add("tDigest", list(range(1, 10)))
    assert 0.1 == round((await modclient.tdigest().cdf("tDigest", 1.0))[0], 1)
    assert 0.9 == round((await modclient.tdigest().cdf("tDigest", 9.0))[0], 1)
    res = await modclient.tdigest().cdf("tDigest", 1.0, 9.0)
    assert [0.1, 0.9] == [round(x, 1) for x in res]


@pytest.mark.redismod
@pytest.mark.experimental
@skip_ifmodversion_lt("2.4.0", "bf")
async def test_tdigest_trimmed_mean(modclient: redis.Redis):
    assert await modclient.tdigest().create("tDigest", 100)
    # insert data-points into sketch
    assert await modclient.tdigest().add("tDigest", list(range(1, 10)))
    assert 5 == await modclient.tdigest().trimmed_mean("tDigest", 0.1, 0.9)
    assert 4.5 == await modclient.tdigest().trimmed_mean("tDigest", 0.4, 0.5)


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_rank(modclient: redis.Redis):
    assert await modclient.tdigest().create("t-digest", 500)
    assert await modclient.tdigest().add("t-digest", list(range(0, 20)))
    assert -1 == (await modclient.tdigest().rank("t-digest", -1))[0]
    assert 0 == (await modclient.tdigest().rank("t-digest", 0))[0]
    assert 10 == (await modclient.tdigest().rank("t-digest", 10))[0]
    assert [-1, 20, 9] == await modclient.tdigest().rank("t-digest", -20, 20, 9)


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_revrank(modclient: redis.Redis):
    assert await modclient.tdigest().create("t-digest", 500)
    assert await modclient.tdigest().add("t-digest", list(range(0, 20)))
    assert -1 == (await modclient.tdigest().revrank("t-digest", 20))[0]
    assert 19 == (await modclient.tdigest().revrank("t-digest", 0))[0]
    assert [-1, 19, 9] == await modclient.tdigest().revrank("t-digest", 21, 0, 10)


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_byrank(modclient: redis.Redis):
    assert await modclient.tdigest().create("t-digest", 500)
    assert await modclient.tdigest().add("t-digest", list(range(1, 11)))
    assert 1 == (await modclient.tdigest().byrank("t-digest", 0))[0]
    assert 10 == (await modclient.tdigest().byrank("t-digest", 9))[0]
    assert (await modclient.tdigest().byrank("t-digest", 100))[0] == inf
    with pytest.raises(redis.ResponseError):
        (await modclient.tdigest().byrank("t-digest", -1))[0]


@pytest.mark.redismod
@pytest.mark.experimental
async def test_tdigest_byrevrank(modclient: redis.Redis):
    assert await modclient.tdigest().create("t-digest", 500)
    assert await modclient.tdigest().add("t-digest", list(range(1, 11)))
    assert 10 == (await modclient.tdigest().byrevrank("t-digest", 0))[0]
    assert 1 == (await modclient.tdigest().byrevrank("t-digest", 9))[0]
    assert (await modclient.tdigest().byrevrank("t-digest", 100))[0] == -inf
    with pytest.raises(redis.ResponseError):
        (await modclient.tdigest().byrevrank("t-digest", -1))[0]


# @pytest.mark.redismod
# async def test_pipeline(modclient: redis.Redis):
#     pipeline = await modclient.bf().pipeline()
#     assert not await modclient.bf().execute_command("get pipeline")
#
#     assert await modclient.bf().create("pipeline", 0.01, 1000)
#     for i in range(100):
#         pipeline.add("pipeline", i)
#     for i in range(100):
#         assert not (await modclient.bf().exists("pipeline", i))
#
#     pipeline.execute()
#
#     for i in range(100):
#         assert await modclient.bf().exists("pipeline", i)