This is an automated email from the ASF dual-hosted git repository.
nic-6443 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/apisix.git
The following commit(s) were added to refs/heads/master by this push:
new 5e3ad1f107 test(ai-rate-limiting): wait for token usage to reach redis
between requests (#13841)
5e3ad1f107 is described below
commit 5e3ad1f10782e4790124d73af16092c70e9de7e2
Author: Nic <[email protected]>
AuthorDate: Wed Aug 19 08:52:48 2026 +0800
test(ai-rate-limiting): wait for token usage to reach redis between
requests (#13841)
---
t/lib/test_redis.lua | 117 +++++++++++++++++++++++++++++++++++---------
t/plugin/ai-rate-limiting.t | 54 +++++++++++++++-----
2 files changed, 137 insertions(+), 34 deletions(-)
diff --git a/t/lib/test_redis.lua b/t/lib/test_redis.lua
index 7f097a4608..6947bb01c9 100644
--- a/t/lib/test_redis.lua
+++ b/t/lib/test_redis.lua
@@ -70,7 +70,26 @@ local function auth_if_needed(red, opts)
return nil, err
end
-local function flush_single(host, port, opts)
+local function release(red, host, port, opts)
+ local keepalive_pool = opts.keepalive_pool
+ if keepalive_pool == nil then
+ keepalive_pool = 0
+ end
+ if keepalive_pool == 0 then
+ local ok_close, close_err = red:close()
+ if not ok_close then
+ log_warn("failed to close redis connection ", host, ":", port, ":
", close_err)
+ end
+ else
+ local keepalive_timeout = opts.keepalive_timeout or 10000
+ local ok_keepalive, keepalive_err =
red:set_keepalive(keepalive_timeout, keepalive_pool)
+ if not ok_keepalive then
+ log_warn("failed to set keepalive for redis ", host, ":", port, ":
", keepalive_err)
+ end
+ end
+end
+
+local function connect(host, port, opts)
local red = redis:new()
local connect_timeout = opts.connect_timeout or 1000
local send_timeout = opts.send_timeout or connect_timeout
@@ -85,35 +104,34 @@ local function flush_single(host, port, opts)
local ok_auth, auth_err = auth_if_needed(red, opts)
if not ok_auth then
- local ok_close, close_err = red:close()
- if not ok_close then
- log_warn("failed to close redis connection ", host, ":", port, ":
", close_err)
- end
+ release(red, host, port, {})
return nil, auth_err
end
+ if opts.database then
+ local ok_select, select_err = red:select(opts.database)
+ if not ok_select then
+ log_warn("failed to select redis db ", opts.database, ": ",
select_err)
+ release(red, host, port, {})
+ return nil, select_err
+ end
+ end
+
+ return red
+end
+
+local function flush_single(host, port, opts)
+ local red, err = connect(host, port, opts)
+ if not red then
+ return nil, err
+ end
+
local _, flush_err = red:flushall()
if flush_err then
log_warn("failed to flush redis ", host, ":", port, ": ", flush_err)
end
- local keepalive_pool = opts.keepalive_pool
- if keepalive_pool == nil then
- keepalive_pool = 0
- end
- if keepalive_pool == 0 then
- local ok_close, close_err = red:close()
- if not ok_close then
- log_warn("failed to close redis connection ", host, ":", port, ":
", close_err)
- end
- else
- local keepalive_timeout = opts.keepalive_timeout or 10000
- keepalive_pool = keepalive_pool or 100
- local ok_keepalive, keepalive_err =
red:set_keepalive(keepalive_timeout, keepalive_pool)
- if not ok_keepalive then
- log_warn("failed to set keepalive for redis ", host, ":", port, ":
", keepalive_err)
- end
- end
+ release(red, host, port, opts)
return true
end
@@ -160,6 +178,61 @@ function _M.flush_port(host, port, opts)
return flush_single(host, port, opts)
end
+-- Rate limit counters are committed from a log phase timer, so a request fired
+-- right after the previous one can still be judged against a budget that looks
+-- unspent. Snapshot the counters with sum_counters() before a request and wait
+-- for the total to grow with wait_counters_above() afterwards, rather than
+-- sleeping a fixed amount: a wait long enough to cover a busy machine would
roll
+-- over the short time windows some of these tests configure.
+function _M.sum_counters(pattern, opts)
+ opts = opts or {}
+ local host = opts.host or DEFAULT_HOST
+ local port = opts.port or 6379
+
+ local red, err = connect(host, port, opts)
+ if not red then
+ return nil, err
+ end
+
+ local keys, keys_err = red:keys(pattern)
+ if not keys then
+ release(red, host, port, opts)
+ return nil, keys_err
+ end
+
+ local total = 0
+ for _, key in ipairs(keys) do
+ -- a key that expired between keys() and get() reads back as ngx.null,
+ -- which is not an error; keep the two apart so a failed read cannot be
+ -- mistaken for a counter of zero
+ local value, get_err = red:get(key)
+ if not value then
+ release(red, host, port, opts)
+ return nil, get_err
+ end
+ total = total + (tonumber(value) or 0)
+ end
+
+ release(red, host, port, opts)
+
+ return total
+end
+
+function _M.wait_counters_above(pattern, previous, opts)
+ local last_err
+ for _ = 1, 100 do
+ local total, err = _M.sum_counters(pattern, opts)
+ if total and total > previous then
+ return true
+ end
+ last_err = err
+ ngx.sleep(0.01)
+ end
+
+ return nil, "counters matching " .. pattern .. " stayed at " .. previous ..
+ (last_err and (", last error: " .. last_err) or "")
+end
+
_M.default_ports = DEFAULT_PORTS
return _M
diff --git a/t/plugin/ai-rate-limiting.t b/t/plugin/ai-rate-limiting.t
index 9fcadbc1c7..d4fad74991 100644
--- a/t/plugin/ai-rate-limiting.t
+++ b/t/plugin/ai-rate-limiting.t
@@ -1630,18 +1630,48 @@ passed
=== TEST 37: redis policy shares counter and rejects the 4th request
---- pipelined_requests eval
-[
- "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\":
\"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is
1+1?\"} ] }",
- "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\":
\"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is
1+1?\"} ] }",
- "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\":
\"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is
1+1?\"} ] }",
- "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\":
\"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is
1+1?\"} ] }",
-]
---- more_headers
-Authorization: Bearer token
-X-AI-Fixture: openai/chat-model-echo.json
---- error_code eval
-[200, 200, 200, 503]
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local test_redis = require("lib.test_redis")
+ local COUNTERS = "plugin-ai-rate-limiting:*"
+ local OPTS = {database = 1}
+
+ local httpc = http.new()
+ local codes = {}
+ for i = 1, 4 do
+ local before, before_err = test_redis.sum_counters(COUNTERS,
OPTS)
+ assert(before, before_err)
+ local res = assert(httpc:request_uri(
+ "http://127.0.0.1:"; .. ngx.var.server_port .. "/ai",
+ {
+ method = "POST",
+ body = [[{
+ "messages": [
+ { "role": "system", "content": "You are a
mathematician" },
+ { "role": "user", "content": "What is 1+1?" }
+ ]
+ }]],
+ headers = {
+ ["Content-Type"] = "application/json",
+ ["Authorization"] = "Bearer token",
+ ["X-AI-Fixture"] = "openai/chat-model-echo.json",
+ }
+ }
+ ))
+ codes[i] = res.status
+ -- only a request that reached the LLM has usage to commit
+ if res.status == 200 and i < 4 then
+ assert(test_redis.wait_counters_above(COUNTERS, before,
OPTS))
+ end
+ end
+ ngx.say(table.concat(codes, ", "))
+ }
+ }
+--- timeout: 10
+--- response_body
+200, 200, 200, 503
