From a0828f6ddf755d78dd2151513b49348ebfc85a94 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Thu, 21 Sep 2023 22:54:48 +0300
Subject: [PATCH] #1317 Fixed priority of events isolation (#1318)

---
 CHANGES/1317.bugfix.rst                       |  1 +
 aiogram/fsm/middleware.py                     |  4 +-
 .../test_1317_state_vs_isolation.py           | 81 +++++++++++++++++++
 3 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 CHANGES/1317.bugfix.rst
 create mode 100644 tests/test_issues/test_1317_state_vs_isolation.py

diff --git a/CHANGES/1317.bugfix.rst b/CHANGES/1317.bugfix.rst
new file mode 100644
index 00000000..dcdf1e47
--- /dev/null
+++ b/CHANGES/1317.bugfix.rst
@@ -0,0 +1 @@
+Fixed priority of events isolation, now user state will be loaded only after lock is acquired
diff --git a/aiogram/fsm/middleware.py b/aiogram/fsm/middleware.py
index 6de91a83..d1f1d973 100644
--- a/aiogram/fsm/middleware.py
+++ b/aiogram/fsm/middleware.py
@@ -34,8 +34,10 @@ class FSMContextMiddleware(BaseMiddleware):
         context = self.resolve_event_context(bot, data)
         data["fsm_storage"] = self.storage
         if context:
-            data.update({"state": context, "raw_state": await context.get_state()})
+            # Bugfix: https://github.com/aiogram/aiogram/issues/1317
+            # State should be loaded after lock is acquired
             async with self.events_isolation.lock(key=context.key):
+                data.update({"state": context, "raw_state": await context.get_state()})
                 return await handler(event, data)
         return await handler(event, data)
 
diff --git a/tests/test_issues/test_1317_state_vs_isolation.py b/tests/test_issues/test_1317_state_vs_isolation.py
new file mode 100644
index 00000000..31613fa4
--- /dev/null
+++ b/tests/test_issues/test_1317_state_vs_isolation.py
@@ -0,0 +1,81 @@
+import asyncio
+from datetime import datetime
+
+from aiogram import Dispatcher
+from aiogram.filters import Command
+from aiogram.fsm.context import FSMContext
+from aiogram.fsm.state import State, StatesGroup
+from aiogram.fsm.storage.memory import SimpleEventIsolation
+from aiogram.types import Chat, Message, Update, User
+from tests.mocked_bot import MockedBot
+
+
+class TestStateVSIsolation:
+    async def test_issue(self, bot: MockedBot):
+        dispatcher = Dispatcher(events_isolation=SimpleEventIsolation())
+        first = 0
+        second = 0
+        third = 0
+        stack = []
+
+        class TestState(StatesGroup):
+            foo = State()
+            bar = State()
+            baz = State()
+
+        @dispatcher.message(Command("test"))
+        async def command_top(message: Message, state: FSMContext):
+            nonlocal first
+            first += 1
+            stack.append("command")
+            await state.set_state(TestState.foo)
+
+        @dispatcher.message(TestState.foo)
+        async def handle_foo(message: Message, state: FSMContext):
+            nonlocal second
+            second += 1
+            stack.append("foo")
+            await state.set_state(TestState.bar)
+
+        @dispatcher.message(TestState.bar)
+        async def handle_bar(message: Message, state: FSMContext):
+            nonlocal third
+            third += 1
+            stack.append("bar")
+            await state.set_state(None)
+
+        @dispatcher.message()
+        async def handle_all(message: Message):
+            stack.append("all")
+
+        await asyncio.gather(
+            *(
+                dispatcher.feed_update(bot, update)
+                for update in [
+                    create_message_update(index=1, text="/test"),
+                    create_message_update(index=2, text="foo"),
+                    create_message_update(index=3, text="bar"),
+                    create_message_update(index=4, text="baz"),
+                ]
+            )
+        )
+
+        # Before bug fix:
+        #   first == 1, second == 3, third == 0, stack == ["command", "foo", "foo", "foo"]
+        assert first == 1
+        assert second == 1
+        assert third == 1
+        assert stack == ["command", "foo", "bar", "all"]
+
+
+def create_message_update(index: int, text: str):
+    return Update(
+        update_id=index,
+        message=Message(
+            message_id=index,
+            date=datetime.now(),
+            chat=Chat(id=42, type="private"),
+            from_user=User(id=42, is_bot=False, first_name="Test", username="test"),
+            text=text,
+        ),
+    )