| 1 | import sys |
| 2 | import threading |
| 3 | import unittest |
| 4 | from pathlib import Path |
| 5 | |
| 6 | sys.path.insert(0, str(Path(__file__).parent.parent.parent)) |
| 7 | |
| 8 | from app.models import const |
| 9 | from app.services.state import MemoryState, RedisState |
| 10 | |
| 11 | |
| 12 | class _FakeRedis: |
| 13 | def __init__(self, batches): |
| 14 | self.batches = batches |
| 15 | self.data = {} |
| 16 | for key in [key for batch in batches for key in batch]: |
| 17 | index = int(key.decode("utf-8").split(":")[-1]) |
| 18 | self.data[key] = { |
| 19 | b"task_id": key, |
| 20 | b"state": b"1", |
| 21 | b"progress": str(index).encode("utf-8"), |
| 22 | } |
| 23 | |
| 24 | def scan(self, cursor, count): |
| 25 | batch_index = int(cursor) |
| 26 | next_cursor = batch_index + 1 |
| 27 | if next_cursor >= len(self.batches): |
| 28 | next_cursor = 0 |
| 29 | return next_cursor, self.batches[batch_index] |
| 30 | |
| 31 | def hgetall(self, key): |
| 32 | return self.data[key] |
| 33 | |
| 34 | |
| 35 | class TestMemoryState(unittest.TestCase): |
| 36 | def test_get_task_and_get_all_tasks_return_isolated_snapshots(self): |
| 37 | state = MemoryState() |
| 38 | state.update_task( |
| 39 | "task-1", |
| 40 | state=const.TASK_STATE_PROCESSING, |
| 41 | progress=25, |
| 42 | videos=["first.mp4"], |
| 43 | ) |
| 44 | |
| 45 | task = state.get_task("task-1") |
| 46 | task["videos"].append("mutated.mp4") |
| 47 | |
| 48 | tasks, total = state.get_all_tasks(page=1, page_size=10) |
| 49 | tasks[0]["videos"].append("mutated-again.mp4") |
| 50 | |
| 51 | self.assertEqual(total, 1) |
| 52 | self.assertEqual(state.get_task("task-1")["videos"], ["first.mp4"]) |
| 53 | |
| 54 | def test_concurrent_memory_updates_are_preserved(self): |
| 55 | state = MemoryState() |
| 56 | thread_count = 5 |
| 57 | tasks_per_thread = 50 |
| 58 | |
| 59 | def update_tasks(thread_index): |
| 60 | for task_index in range(tasks_per_thread): |
| 61 | state.update_task( |
| 62 | f"task-{thread_index}-{task_index}", |
| 63 | state=const.TASK_STATE_PROCESSING, |
| 64 | progress=task_index, |
| 65 | ) |
| 66 | |
| 67 | threads = [ |
| 68 | threading.Thread(target=update_tasks, args=(thread_index,)) |
| 69 | for thread_index in range(thread_count) |
| 70 | ] |
| 71 | for thread in threads: |
| 72 | thread.start() |
| 73 | for thread in threads: |
| 74 | thread.join() |
| 75 | |
| 76 | tasks, total = state.get_all_tasks(page=1, page_size=thread_count * tasks_per_thread) |
| 77 | |
| 78 | self.assertEqual(total, thread_count * tasks_per_thread) |
| 79 | self.assertEqual(len(tasks), total) |
| 80 | |
| 81 | |
| 82 | class TestRedisState(unittest.TestCase): |
| 83 | def _build_state(self, batch_sizes): |
| 84 | keys = [f"task:{i}".encode("utf-8") for i in range(sum(batch_sizes))] |
| 85 | batches = [] |
| 86 | offset = 0 |
| 87 | for batch_size in batch_sizes: |
| 88 | batches.append(keys[offset : offset + batch_size]) |
| 89 | offset += batch_size |
| 90 | |
| 91 | state = RedisState.__new__(RedisState) |
| 92 | state._redis = _FakeRedis(batches) |
| 93 | return state |
| 94 | |
| 95 | def test_get_all_tasks_paginates_across_scan_batches(self): |
| 96 | """ |
| 97 | Redis SCAN 分批返回 key 时，分页切片必须按当前批次起始位置计算。 |
| 98 | |
| 99 | 这个用例复现 PR #890 描述的 18 条任务、page_size=10 场景： |
| 100 | 第一批 10 条，第二批 8 条。旧逻辑第一页会返回空列表，第二页 |
| 101 | 只返回 2 条；修复后第一页返回 10 条，第二页返回剩余 8 条。 |
| 102 | """ |
| 103 | state = self._build_state([10, 8]) |
| 104 | |
| 105 | first_page, first_total = state.get_all_tasks(page=1, page_size=10) |
| 106 | second_page, second_total = state.get_all_tasks(page=2, page_size=10) |
| 107 | |
| 108 | self.assertEqual(first_total, 18) |
| 109 | self.assertEqual(second_total, 18) |
| 110 | self.assertEqual(len(first_page), 10) |
| 111 | self.assertEqual(len(second_page), 8) |
| 112 | self.assertEqual( |
| 113 | [task["task_id"] for task in first_page], |
| 114 | [f"task:{i}" for i in range(10)], |
| 115 | ) |
| 116 | self.assertEqual( |
| 117 | [task["task_id"] for task in second_page], |
| 118 | [f"task:{i}" for i in range(10, 18)], |
| 119 | ) |
| 120 | |
| 121 | |
| 122 | if __name__ == "__main__": |
| 123 | unittest.main() |
| 124 |