From 44273e17ea07c91e37e10357bf6eaa212b25d537 Mon Sep 17 00:00:00 2001
From: kievzenit <kir.backender@gmail.com>
Date: Mon, 11 Aug 2025 17:36:30 +0300
Subject: [PATCH] added more tests to pymongo test, to check for all possible
 cases of using update_data method

---
 tests/test_fsm/storage/test_pymongo.py | 41 ++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/tests/test_fsm/storage/test_pymongo.py b/tests/test_fsm/storage/test_pymongo.py
index 7b87e589..81ef1cdf 100644
--- a/tests/test_fsm/storage/test_pymongo.py
+++ b/tests/test_fsm/storage/test_pymongo.py
@@ -33,6 +33,47 @@ async def test_update_not_existing_data_with_empty_dictionary(
     assert await pymongo_storage._collection.find_one({}) is None
 
 
+async def test_update_not_existing_data_with_non_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    assert await pymongo_storage.update_data(key=storage_key, data={"key": "value"}) == {
+        "key": "value"
+    }
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "data": {"key": "value"},
+    }
+    await pymongo_storage._collection.delete_one({})
+
+
+async def test_update_existing_data_with_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    await pymongo_storage.set_data(key=storage_key, data={"key": "value"})
+    assert await pymongo_storage.update_data(key=storage_key, data={}) == {}
+    assert await pymongo_storage._collection.find_one({}) is None
+
+
+async def test_update_existing_data_with_non_empty_dictionary(
+    pymongo_storage: PyMongoStorage,
+    storage_key: StorageKey,
+):
+    assert await pymongo_storage._collection.find_one({}) is None
+    await pymongo_storage.set_data(key=storage_key, data={"key": "value"})
+    assert await pymongo_storage.update_data(key=storage_key, data={"key": "new_value"}) == {
+        "key": "new_value"
+    }
+    assert await pymongo_storage._collection.find_one({}) == {
+        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
+        "data": {"key": "new_value"},
+    }
+    await pymongo_storage._collection.delete_one({})
+
+
 async def test_document_life_cycle(
     pymongo_storage: PyMongoStorage,
     storage_key: StorageKey,