aiogram/tests/test_fsm/storage/test_mongo.py

import pytest
from pymongo.errors import PyMongoError

from aiogram.fsm.state import State
from aiogram.fsm.storage.mongo import MongoStorage, StorageKey
from tests.conftest import CHAT_ID, USER_ID

PREFIX = "fsm"


async def test_get_storage_passing_only_url(mongo_server):
    storage = MongoStorage.from_url(url=mongo_server)
    try:
        await storage._client.server_info()
    except PyMongoError as e:
        pytest.fail(str(e))


async def test_update_not_existing_data_with_empty_dictionary(
    mongo_storage: MongoStorage,
    storage_key: StorageKey,
):
    assert await mongo_storage._collection.find_one({}) is None
    assert await mongo_storage.get_data(key=storage_key) == {}
    assert await mongo_storage.update_data(key=storage_key, data={}) == {}
    assert await mongo_storage._collection.find_one({}) is None


async def test_document_life_cycle(
    mongo_storage: MongoStorage,
    storage_key: StorageKey,
):
    assert await mongo_storage._collection.find_one({}) is None
    await mongo_storage.set_state(storage_key, "test")
    await mongo_storage.set_data(storage_key, {"key": "value"})
    assert await mongo_storage._collection.find_one({}) == {
        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
        "state": "test",
        "data": {"key": "value"},
    }
    await mongo_storage.set_state(storage_key, None)
    assert await mongo_storage._collection.find_one({}) == {
        "_id": f"{PREFIX}:{CHAT_ID}:{USER_ID}",
        "data": {"key": "value"},
    }
    await mongo_storage.set_data(storage_key, {})
    assert await mongo_storage._collection.find_one({}) is None


class TestStateAndDataDoNotAffectEachOther:
    async def test_state_and_data_do_not_affect_each_other_while_getting(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        assert await mongo_storage._collection.find_one({}) is None
        await mongo_storage.set_state(storage_key, "test")
        await mongo_storage.set_data(storage_key, {"key": "value"})
        assert await mongo_storage.get_state(storage_key) == "test"
        assert await mongo_storage.get_data(storage_key) == {"key": "value"}

    async def test_data_do_not_affect_to_deleted_state_getting(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        await mongo_storage.set_state(storage_key, "test")
        await mongo_storage.set_data(storage_key, {"key": "value"})
        await mongo_storage.set_state(storage_key, None)
        assert await mongo_storage.get_state(storage_key) is None

    async def test_state_do_not_affect_to_deleted_data_getting(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        await mongo_storage.set_state(storage_key, "test")
        await mongo_storage.set_data(storage_key, {"key": "value"})
        await mongo_storage.set_data(storage_key, {})
        assert await mongo_storage.get_data(storage_key) == {}

    async def test_state_do_not_affect_to_updating_not_existing_data_with_empty_dictionary(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        await mongo_storage.set_state(storage_key, "test")
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test"
        }
        assert await mongo_storage.update_data(key=storage_key, data={}) == {}
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test"
        }

    async def test_state_do_not_affect_to_updating_not_existing_data_with_non_empty_dictionary(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        await mongo_storage.set_state(storage_key, "test")
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test"
        }
        assert await mongo_storage.update_data(
            key=storage_key,
            data={"key": "value"},
        ) == {"key": "value"}
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test",
            "data": {"key": "value"},
        }

    async def test_state_do_not_affect_to_updating_existing_data_with_empty_dictionary(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        await mongo_storage.set_state(storage_key, "test")
        await mongo_storage.set_data(storage_key, {"key": "value"})
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test",
            "data": {"key": "value"},
        }
        assert await mongo_storage.update_data(key=storage_key, data={}) == {"key": "value"}
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test",
            "data": {"key": "value"},
        }

    async def test_state_do_not_affect_to_updating_existing_data_with_non_empty_dictionary(
        self,
        mongo_storage: MongoStorage,
        storage_key: StorageKey,
    ):
        await mongo_storage.set_state(storage_key, "test")
        await mongo_storage.set_data(storage_key, {"key": "value"})
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test",
            "data": {"key": "value"},
        }
        assert await mongo_storage.update_data(
            key=storage_key,
            data={"key": "VALUE", "key_2": "value_2"},
        ) == {"key": "VALUE", "key_2": "value_2"}
        assert await mongo_storage._collection.find_one({}, projection={"_id": 0}) == {
            "state": "test",
            "data": {"key": "VALUE", "key_2": "value_2"},
        }


@pytest.mark.parametrize(
    "value,result",
    [
        [None, None],
        ["", ""],
        ["text", "text"],
        [State(), None],
        [State(state="*"), "*"],
        [State("text"), "@:text"],
        [State("test", group_name="Test"), "Test:test"],
        [[1, 2, 3], "[1, 2, 3]"],
    ],
)
def test_resolve_state(value, result, mongo_storage: MongoStorage):
    assert mongo_storage.resolve_state(value) == result