From 03ccebd8be7b4ba463925a0516ab1ed7fa147fd2 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Thu, 13 May 2021 01:54:07 +0300
Subject: [PATCH] Cover base and memory storage

---
 aiogram/dispatcher/fsm/storage/base.py        | 14 +++++---
 tests/test_dispatcher/fsm/storage/__init__.py |  0
 .../test_dispatcher/fsm/storage/test_base.py  | 34 +++++++++++++++++++
 3 files changed, 43 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_dispatcher/fsm/storage/__init__.py
 create mode 100644 tests/test_dispatcher/fsm/storage/test_base.py

diff --git a/aiogram/dispatcher/fsm/storage/base.py b/aiogram/dispatcher/fsm/storage/base.py
index 29cce9ed..36cebb31 100644
--- a/aiogram/dispatcher/fsm/storage/base.py
+++ b/aiogram/dispatcher/fsm/storage/base.py
@@ -10,23 +10,27 @@ StateType = Optional[Union[str, State]]
 class BaseStorage(ABC):
     @abstractmethod
     @asynccontextmanager
-    async def lock(self) -> AsyncGenerator[None, None]:
+    async def lock(self) -> AsyncGenerator[None, None]:  # pragma: no cover
         yield None
 
     @abstractmethod
-    async def set_state(self, chat_id: int, user_id: int, state: StateType = None) -> None:
+    async def set_state(
+        self, chat_id: int, user_id: int, state: StateType = None
+    ) -> None:  # pragma: no cover
         pass
 
     @abstractmethod
-    async def get_state(self, chat_id: int, user_id: int) -> Optional[str]:
+    async def get_state(self, chat_id: int, user_id: int) -> Optional[str]:  # pragma: no cover
         pass
 
     @abstractmethod
-    async def set_data(self, chat_id: int, user_id: int, data: Dict[str, Any]) -> None:
+    async def set_data(
+        self, chat_id: int, user_id: int, data: Dict[str, Any]
+    ) -> None:  # pragma: no cover
         pass
 
     @abstractmethod
-    async def get_data(self, chat_id: int, user_id: int) -> Dict[str, Any]:
+    async def get_data(self, chat_id: int, user_id: int) -> Dict[str, Any]:  # pragma: no cover
         pass
 
     async def update_data(
diff --git a/tests/test_dispatcher/fsm/storage/__init__.py b/tests/test_dispatcher/fsm/storage/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_dispatcher/fsm/storage/test_base.py b/tests/test_dispatcher/fsm/storage/test_base.py
new file mode 100644
index 00000000..8b7129a4
--- /dev/null
+++ b/tests/test_dispatcher/fsm/storage/test_base.py
@@ -0,0 +1,34 @@
+import pytest
+
+from aiogram.dispatcher.fsm.storage.memory import MemoryStorage, MemoryStorageRecord
+
+
+@pytest.fixture()
+def storage():
+    return MemoryStorage()
+
+
+class TestMemoryStorage:
+    @pytest.mark.asyncio
+    async def test_set_state(self, storage: MemoryStorage):
+        assert await storage.get_state(chat_id=-42, user_id=42) is None
+
+        await storage.set_state(chat_id=-42, user_id=42, state="state")
+        assert await storage.get_state(chat_id=-42, user_id=42) == "state"
+
+        assert -42 in storage.storage
+        assert 42 in storage.storage[-42]
+        assert isinstance(storage.storage[-42][42], MemoryStorageRecord)
+        assert storage.storage[-42][42].state == "state"
+
+    @pytest.mark.asyncio
+    async def test_set_data(self, storage: MemoryStorage):
+        assert await storage.get_data(chat_id=-42, user_id=42) == {}
+
+        await storage.set_data(chat_id=-42, user_id=42, data={"foo": "bar"})
+        assert await storage.get_data(chat_id=-42, user_id=42) == {"foo": "bar"}
+
+        assert -42 in storage.storage
+        assert 42 in storage.storage[-42]
+        assert isinstance(storage.storage[-42][42], MemoryStorageRecord)
+        assert storage.storage[-42][42].data == {"foo": "bar"}