75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
|
|
|
from app.services.state import RedisState
|
|
|
|
|
|
class _FakeRedis:
|
|
def __init__(self, batches):
|
|
self.batches = batches
|
|
self.data = {}
|
|
for key in [key for batch in batches for key in batch]:
|
|
index = int(key.decode("utf-8").split(":")[-1])
|
|
self.data[key] = {
|
|
b"task_id": key,
|
|
b"state": b"1",
|
|
b"progress": str(index).encode("utf-8"),
|
|
}
|
|
|
|
def scan(self, cursor, count):
|
|
batch_index = int(cursor)
|
|
next_cursor = batch_index + 1
|
|
if next_cursor >= len(self.batches):
|
|
next_cursor = 0
|
|
return next_cursor, self.batches[batch_index]
|
|
|
|
def hgetall(self, key):
|
|
return self.data[key]
|
|
|
|
|
|
class TestRedisState(unittest.TestCase):
|
|
def _build_state(self, batch_sizes):
|
|
keys = [f"task:{i}".encode("utf-8") for i in range(sum(batch_sizes))]
|
|
batches = []
|
|
offset = 0
|
|
for batch_size in batch_sizes:
|
|
batches.append(keys[offset : offset + batch_size])
|
|
offset += batch_size
|
|
|
|
state = RedisState.__new__(RedisState)
|
|
state._redis = _FakeRedis(batches)
|
|
return state
|
|
|
|
def test_get_all_tasks_paginates_across_scan_batches(self):
|
|
"""
|
|
Redis SCAN 分批返回 key 时，分页切片必须按当前批次起始位置计算。
|
|
|
|
这个用例复现 PR #890 描述的 18 条任务、page_size=10 场景：
|
|
第一批 10 条，第二批 8 条。旧逻辑第一页会返回空列表，第二页
|
|
只返回 2 条；修复后第一页返回 10 条，第二页返回剩余 8 条。
|
|
"""
|
|
state = self._build_state([10, 8])
|
|
|
|
first_page, first_total = state.get_all_tasks(page=1, page_size=10)
|
|
second_page, second_total = state.get_all_tasks(page=2, page_size=10)
|
|
|
|
self.assertEqual(first_total, 18)
|
|
self.assertEqual(second_total, 18)
|
|
self.assertEqual(len(first_page), 10)
|
|
self.assertEqual(len(second_page), 8)
|
|
self.assertEqual(
|
|
[task["task_id"] for task in first_page],
|
|
[f"task:{i}" for i in range(10)],
|
|
)
|
|
self.assertEqual(
|
|
[task["task_id"] for task in second_page],
|
|
[f"task:{i}" for i in range(10, 18)],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|