Small changes in tests

aiogram · Jun 14, 2021 · 5296724 · 5296724
1 parent 7f0a9d8
commit 5296724
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -63,6 +63,8 @@ async def memory_storage():
 def bot():
     bot = MockedBot()
     token = Bot.set_current(bot)
-    yield bot
-    Bot.reset_current(token)
-    bot.me.invalidate(bot)
+    try:
+        yield bot
+    finally:
+        Bot.reset_current(token)
+        bot.me.invalidate(bot)
diff --git a/tests/test_dispatcher/test_fsm/storage/test_storages.py b/tests/test_dispatcher/test_fsm/storage/test_storages.py
@@ -42,3 +42,7 @@ async def test_update_data(self, bot: MockedBot, storage: BaseStorage):
         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) == {
+            "foo": "bar",
+            "baz": "spam",
+        }