dragonfly/tests/fakeredis/test/test_lua_modules.py

import json

import pytest
import redis

pytestmark = []
pytestmark.extend(
    [
        pytest.mark.asyncio,
    ]
)

lua_modules_test = pytest.importorskip("lupa")


@pytest.mark.load_lua_modules("cjson")
async def test_async_asgi_ratelimit_script(async_redis: redis.Redis):
    script = """
local ruleset = cjson.decode(ARGV[1])

-- Set limits
for i, key in pairs(KEYS) do
    redis.call('SET', key, ruleset[key][1], 'EX', ruleset[key][2], 'NX')
end

-- Check limits
for i = 1, #KEYS do
    local value = redis.call('GET', KEYS[i])
    if value and tonumber(value) < 1 then
        return ruleset[KEYS[i]][2]
    end
end

-- Decrease limits
for i, key in pairs(KEYS) do
    redis.call('DECR', key)
end
return 0
"""

    script = async_redis.register_script(script)
    ruleset = {"path:get:user:name": (1, 1)}
    await script(keys=list(ruleset.keys()), args=[json.dumps(ruleset)])


@pytest.mark.load_lua_modules("cjson")
def test_asgi_ratelimit_script(r: redis.Redis):
    script = """
local ruleset = cjson.decode(ARGV[1])

-- Set limits
for i, key in pairs(KEYS) do
    redis.call('SET', key, ruleset[key][1], 'EX', ruleset[key][2], 'NX')
end

-- Check limits
for i = 1, #KEYS do
    local value = redis.call('GET', KEYS[i])
    if value and tonumber(value) < 1 then
        return ruleset[KEYS[i]][2]
    end
end

-- Decrease limits
for i, key in pairs(KEYS) do
    redis.call('DECR', key)
end
return 0
"""

    script = r.register_script(script)
    ruleset = {"path:get:user:name": (1, 1)}
    script(keys=list(ruleset.keys()), args=[json.dumps(ruleset)])