From 47d6456ebc42876b5556929c6968a543ae06545b Mon Sep 17 00:00:00 2001 From: yuetsh <517252939@qq.com> Date: Thu, 7 May 2026 19:56:16 -0600 Subject: [PATCH] fix cache --- utils/cache.py | 9 +++++++++ utils/tests.py | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) create mode 100644 utils/tests.py diff --git a/utils/cache.py b/utils/cache.py index 3d5e74f..fa39bdb 100644 --- a/utils/cache.py +++ b/utils/cache.py @@ -11,6 +11,15 @@ class MyRedisClient(DefaultClient): client = self.get_client(write=True) return getattr(client, item) + def hget(self, name, key, version=None, client=None): + if client is None: + client = self.get_client(write=True) + nkey = self.make_key(key, version=version) + value = client.hget(name, nkey) + if value is None: + return None + return self.decode(value) + def redis_incr(self, key, count=1): """ django 默认的 incr 在 key 不存在时候会抛异常 diff --git a/utils/tests.py b/utils/tests.py new file mode 100644 index 0000000..09869d6 --- /dev/null +++ b/utils/tests.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace +from unittest import TestCase + +from django.conf import settings + +if not settings.configured: + settings.configure( + DJANGO_REDIS_CONNECTION_FACTORY="django_redis.pool.ConnectionFactory", + ) + +from utils.cache import MyRedisClient +from utils.throttling import TokenBucket + + +class FakeRedis: + def __init__(self): + self.hashes = {} + + def hset(self, name, key, value): + self.hashes.setdefault(str(name), {})[str(key)] = value + return 1 + + def hget(self, name, key): + return self.hashes.get(str(name), {}).get(str(key)) + + +def make_client(): + backend = SimpleNamespace( + key_prefix="", + version=1, + key_func=lambda key, key_prefix, version: key, + ) + client = MyRedisClient("redis://localhost:6379/0", {}, backend) + fake_redis = FakeRedis() + client.get_client = lambda write=True: fake_redis + return client + + +class TokenBucketCacheTests(TestCase): + def test_token_bucket_reads_values_written_through_django_redis_hash_api(self): + bucket = TokenBucket( + key="submission:user:1", + capacity=2, + fill_rate=0.1, + default_capacity=2, + redis_conn=make_client(), + ) + + self.assertEqual(bucket.consume(), (True, 0)) + self.assertEqual(bucket.consume(), (True, 0))