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 import Bot
from aiogram.dispatcher.fsm.state import State 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} DEFAULT_REDIS_LOCK_KWARGS = {"timeout": 60}
@ -62,6 +62,13 @@ class DefaultKeyBuilder(KeyBuilder):
parts.extend([str(key.chat_id), str(key.user_id)]) parts.extend([str(key.chat_id), str(key.user_id)])
if self.with_destiny: if self.with_destiny:
parts.append(key.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) parts.append(part)
return self.separator.join(parts) return self.separator.join(parts)

View file

@ -1,6 +1,8 @@
from typing import Literal
import pytest 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 from aiogram.dispatcher.fsm.storage.redis import DefaultKeyBuilder
pytestmark = pytest.mark.asyncio pytestmark = pytest.mark.asyncio
@ -10,7 +12,6 @@ BOT_ID = 42
CHAT_ID = -1 CHAT_ID = -1
USER_ID = 2 USER_ID = 2
FIELD = "data" FIELD = "data"
DESTINY = "testing"
class TestRedisDefaultKeyBuilder: class TestRedisDefaultKeyBuilder:
@ -19,8 +20,8 @@ class TestRedisDefaultKeyBuilder:
[ [
[False, False, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"], [False, False, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{FIELD}"],
[True, False, f"{PREFIX}:{BOT_ID}:{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}"], [True, True, f"{PREFIX}:{BOT_ID}:{CHAT_ID}:{USER_ID}:{DEFAULT_DESTINY}:{FIELD}"],
[False, True, f"{PREFIX}:{CHAT_ID}:{USER_ID}:{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): 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_bot_id=with_bot_id,
with_destiny=with_destiny, 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 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)