Check destiny in case of no with_destiny (#776)

* cehck if destiny is default

* Test testiny check

* reformat
This commit is contained in:
Andrey Tikhonov 2021-12-12 18:54:01 +03:00 committed by GitHub
parent 9ec689b562
commit 1357a7a91b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 6 deletions

View file

@ -6,7 +6,7 @@ from aioredis import ConnectionPool, Redis
from aiogram import Bot
from aiogram.dispatcher.fsm.state import State
from aiogram.dispatcher.fsm.storage.base import BaseStorage, StateType, StorageKey
from aiogram.dispatcher.fsm.storage.base import DEFAULT_DESTINY, BaseStorage, StateType, StorageKey
DEFAULT_REDIS_LOCK_KWARGS = {"timeout": 60}
@ -62,6 +62,13 @@ class DefaultKeyBuilder(KeyBuilder):
parts.extend([str(key.chat_id), str(key.user_id)])
if self.with_destiny:
parts.append(key.destiny)
elif key.destiny != DEFAULT_DESTINY:
raise ValueError(
"Redis key builder is not configured to use key destiny other the default.\n"
"\n"
"Probably, you should set `with_destiny=True` in for DefaultKeyBuilder.\n"
"E.g: `RedisStorage(redis, key_builder=DefaultKeyBuilder(with_destiny=True))`"
)
parts.append(part)
return self.separator.join(parts)

View file

@ -1,6 +1,8 @@
from typing import Literal
import pytest
from aiogram.dispatcher.fsm.storage.base import StorageKey
from aiogram.dispatcher.fsm.storage.base import DEFAULT_DESTINY, StorageKey
from aiogram.dispatcher.fsm.storage.redis import DefaultKeyBuilder
pytestmark = pytest.mark.asyncio
@ -10,7 +12,6 @@ BOT_ID = 42
CHAT_ID = -1
USER_ID = 2
FIELD = "data"
DESTINY = "testing"
class TestRedisDefaultKeyBuilder:
@ -19,8 +20,8 @@ class TestRedisDefaultKeyBuilder:
[
[False, False, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"],
[True, False, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{FIELD}"],
[True, True, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DESTINY}:{FIELD}"],
[False, True, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DESTINY}:{FIELD}"],
[True, True, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
[False, True, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
],
)
async def test_generate_key(self, with_bot_id: bool, with_destiny: bool, result: str):
@ -29,5 +30,18 @@ class TestRedisDefaultKeyBuilder:
with_bot_id=with_bot_id,
with_destiny=with_destiny,
)
key = StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=BOT_ID, destiny=DESTINY)
key = StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=BOT_ID, destiny=DEFAULT_DESTINY)
assert key_builder.build(key, FIELD) == result
async def test_destiny_check(self):
key_builder = DefaultKeyBuilder(
with_destiny=False,
)
key = StorageKey(chat_id=CHAT_ID, user_id=USER_ID, bot_id=BOT_ID)
assert key_builder.build(key, FIELD)
key = StorageKey(
chat_id=CHAT_ID, user_id=USER_ID, bot_id=BOT_ID, destiny="CUSTOM_TEST_DESTINY"
)
with pytest.raises(ValueError):
key_builder.build(key, FIELD)