Rework FSM storage key

tests/test_dispatcher/test_fsm/storage/test_redis.py

@@ -1,22 +1,33 @@
 import pytest
 
-from aiogram.dispatcher.fsm.storage.redis import RedisStorage
-from tests.mocked_bot import MockedBot
+from aiogram.dispatcher.fsm.storage.base import StorageKey
+from aiogram.dispatcher.fsm.storage.redis import DefaultKeyBuilder
 
 pytestmark = pytest.mark.asyncio
 
+PREFIX = "test"
+BOT_ID = 42
+CHAT_ID = -1
+USER_ID = 2
+FIELD = "data"
+DESTINY = "testing"
 
-@pytest.mark.redis
-class TestRedisStorage:
+
+class TestRedisDefaultKeyBuilder:
     @pytest.mark.parametrize(
-        "prefix_bot,result",
+        "with_bot_id,with_destiny,result",
         [
-            [False, "fsm:-1:2"],
-            [True, "fsm:42:-1:2"],
-            [{42: "kaboom"}, "fsm:kaboom:-1:2"],
-            [lambda bot: "kaboom", "fsm:kaboom:-1:2"],
+            [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}"],
         ],
     )
-    async def test_generate_key(self, bot: MockedBot, redis_server, prefix_bot, result):
-        storage = RedisStorage.from_url(redis_server, prefix_bot=prefix_bot)
-        assert storage.generate_key(bot, -1, 2) == result
+    async def test_generate_key(self, with_bot_id: bool, with_destiny: bool, result: str):
+        key_builder = DefaultKeyBuilder(
+            prefix=PREFIX,
+            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)
+        assert key_builder.build(key, FIELD) == result

tests/test_dispatcher/test_fsm/storage/test_storages.py

@@ -1,46 +1,54 @@
 import pytest
 
-from aiogram.dispatcher.fsm.storage.base import BaseStorage
+from aiogram.dispatcher.fsm.storage.base import BaseStorage, StorageKey
 from tests.mocked_bot import MockedBot
 
 pytestmark = pytest.mark.asyncio
 
 
+@pytest.fixture(name="storage_key")
+def create_storate_key(bot: MockedBot):
+    return StorageKey(chat_id=-42, user_id=42, bot_id=bot.id)
+
+
 @pytest.mark.parametrize(
     "storage",
     [pytest.lazy_fixture("redis_storage"), pytest.lazy_fixture("memory_storage")],
 )
 class TestStorages:
-    async def test_lock(self, bot: MockedBot, storage: BaseStorage):
+    async def test_lock(self, bot: MockedBot, storage: BaseStorage, storage_key: StorageKey):
         # TODO: ?!?
-        async with storage.lock(bot=bot, chat_id=-42, user_id=42):
+        async with storage.lock(bot=bot, key=storage_key):
             assert True, "You are kidding me?"
 
-    async def test_set_state(self, bot: MockedBot, storage: BaseStorage):
-        assert await storage.get_state(bot=bot, chat_id=-42, user_id=42) is None
+    async def test_set_state(self, bot: MockedBot, storage: BaseStorage, storage_key: StorageKey):
+        assert await storage.get_state(bot=bot, key=storage_key) is None
 
-        await storage.set_state(bot=bot, chat_id=-42, user_id=42, state="state")
-        assert await storage.get_state(bot=bot, chat_id=-42, user_id=42) == "state"
-        await storage.set_state(bot=bot, chat_id=-42, user_id=42, state=None)
-        assert await storage.get_state(bot=bot, chat_id=-42, user_id=42) is None
+        await storage.set_state(bot=bot, key=storage_key, state="state")
+        assert await storage.get_state(bot=bot, key=storage_key) == "state"
+        await storage.set_state(bot=bot, key=storage_key, state=None)
+        assert await storage.get_state(bot=bot, key=storage_key) is None
 
-    async def test_set_data(self, bot: MockedBot, storage: BaseStorage):
-        assert await storage.get_data(bot=bot, chat_id=-42, user_id=42) == {}
+    async def test_set_data(self, bot: MockedBot, storage: BaseStorage, storage_key: StorageKey):
+        assert await storage.get_data(bot=bot, key=storage_key) == {}
 
-        await storage.set_data(bot=bot, chat_id=-42, user_id=42, data={"foo": "bar"})
-        assert await storage.get_data(bot=bot, chat_id=-42, user_id=42) == {"foo": "bar"}
-        await storage.set_data(bot=bot, chat_id=-42, user_id=42, data={})
-        assert await storage.get_data(bot=bot, chat_id=-42, user_id=42) == {}
+        await storage.set_data(bot=bot, key=storage_key, data={"foo": "bar"})
+        assert await storage.get_data(bot=bot, key=storage_key) == {"foo": "bar"}
+        await storage.set_data(bot=bot, key=storage_key, data={})
+        assert await storage.get_data(bot=bot, key=storage_key) == {}
 
-    async def test_update_data(self, bot: MockedBot, storage: BaseStorage):
-        assert await storage.get_data(bot=bot, chat_id=-42, user_id=42) == {}
-        assert await storage.update_data(
-            bot=bot, chat_id=-42, user_id=42, data={"foo": "bar"}
-        ) == {"foo": "bar"}
-        assert await storage.update_data(
-            bot=bot, chat_id=-42, user_id=42, data={"baz": "spam"}
-        ) == {"foo": "bar", "baz": "spam"}
-        assert await storage.get_data(bot=bot, chat_id=-42, user_id=42) == {
+    async def test_update_data(
+        self, bot: MockedBot, storage: BaseStorage, storage_key: StorageKey
+    ):
+        assert await storage.get_data(bot=bot, key=storage_key) == {}
+        assert await storage.update_data(bot=bot, key=storage_key, data={"foo": "bar"}) == {
+            "foo": "bar"
+        }
+        assert await storage.update_data(bot=bot, key=storage_key, data={"baz": "spam"}) == {
+            "foo": "bar",
+            "baz": "spam",
+        }
+        assert await storage.get_data(bot=bot, key=storage_key) == {
             "foo": "bar",
             "baz": "spam",
         }

tests/test_dispatcher/test_fsm/test_context.py

@@ -1,6 +1,7 @@
 import pytest
 
 from aiogram.dispatcher.fsm.context import FSMContext
+from aiogram.dispatcher.fsm.storage.base import StorageKey
 from aiogram.dispatcher.fsm.storage.memory import MemoryStorage
 from tests.mocked_bot import MockedBot
 
@@ -10,21 +11,28 @@ pytestmark = pytest.mark.asyncio
 @pytest.fixture()
 def state(bot: MockedBot):
     storage = MemoryStorage()
-    ctx = storage.storage[bot][-42][42]
+    key = StorageKey(user_id=42, chat_id=-42, bot_id=bot.id)
+    ctx = storage.storage[key]
     ctx.state = "test"
     ctx.data = {"foo": "bar"}
-    return FSMContext(bot=bot, storage=storage, user_id=-42, chat_id=42)
+    return FSMContext(bot=bot, storage=storage, key=key)
 
 
 class TestFSMContext:
     async def test_address_mapping(self, bot: MockedBot):
         storage = MemoryStorage()
-        ctx = storage.storage[bot][-42][42]
+        ctx = storage.storage[StorageKey(chat_id=-42, user_id=42, bot_id=bot.id)]
         ctx.state = "test"
         ctx.data = {"foo": "bar"}
-        state = FSMContext(bot=bot, storage=storage, chat_id=-42, user_id=42)
-        state2 = FSMContext(bot=bot, storage=storage, chat_id=42, user_id=42)
-        state3 = FSMContext(bot=bot, storage=storage, chat_id=69, user_id=69)
+        state = FSMContext(
+            bot=bot, storage=storage, key=StorageKey(chat_id=-42, user_id=42, bot_id=bot.id)
+        )
+        state2 = FSMContext(
+            bot=bot, storage=storage, key=StorageKey(chat_id=42, user_id=42, bot_id=bot.id)
+        )
+        state3 = FSMContext(
+            bot=bot, storage=storage, key=StorageKey(chat_id=69, user_id=69, bot_id=bot.id)
+        )
 
         assert await state.get_state() == "test"
         assert await state2.get_state() is None

tests/test_utils/test_i18n.py

@@ -4,6 +4,7 @@ import pytest
 
 from aiogram import Dispatcher
 from aiogram.dispatcher.fsm.context import FSMContext
+from aiogram.dispatcher.fsm.storage.base import StorageKey
 from aiogram.dispatcher.fsm.storage.memory import MemoryStorage
 from aiogram.types import Update, User
 from aiogram.utils.i18n import ConstI18nMiddleware, FSMI18nMiddleware, I18n, SimpleI18nMiddleware
@@ -131,7 +132,9 @@ class TestFSMI18nMiddleware:
     async def test_middleware(self, i18n: I18n, bot: MockedBot):
         middleware = FSMI18nMiddleware(i18n=i18n)
         storage = MemoryStorage()
-        state = FSMContext(bot=bot, storage=storage, user_id=42, chat_id=42)
+        state = FSMContext(
+            bot=bot, storage=storage, key=StorageKey(user_id=42, chat_id=42, bot_id=bot.id)
+        )
         data = {
             "event_from_user": User(id=42, is_bot=False, language_code="it", first_name="Test"),
             "state": state,