From 5296724a0f77d072634aaf8f8c3f124604b0f506 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Tue, 15 Jun 2021 02:01:57 +0300
Subject: [PATCH] Small changes in tests

---
 tests/conftest.py                                       | 8 +++++---
 tests/test_dispatcher/test_fsm/storage/test_storages.py | 4 ++++
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 92dd97fe..7586f389 100644
--- 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
index fcb2deae..cae3cb2c 100644
--- a/tests/test_dispatcher/test_fsm/storage/test_storages.py
+++ b/tests/test_dispatcher/test_fsm/storage/test_storages.py
@@ -42,3 +42,7 @@ class TestStorages:
         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",
+        }