feat: store data between states

pybotx_fsm/fsm.py

@@ -21,14 +21,18 @@ def __init__(self, state_repo: StateRepoProto, message: IncomingMessage):
         self._message = message
 
     async def change_state(self, state: Enum, **kwargs: Any) -> None:
+        storage_data = await self.get_storage_data()
+        storage_data = vars(storage_data)
+        storage_data.update(**kwargs)
+
         await self._state_repo.set(
             KEY_TEMPLATE.format(
                 host=self._message.bot.host,
                 bot_id=self._message.bot.id,
                 chat_id=self._message.chat.id,
                 user_huid=self._message.sender.huid,
             ),
-            FSMStateData(state, SimpleNamespace(**kwargs)),
+            FSMStateData(state, SimpleNamespace(**storage_data)),
         )
 
     async def get_state(self) -> Optional[Enum]:
@@ -45,6 +49,21 @@ async def get_state(self) -> Optional[Enum]:
             return None
 
         return fsm_state_data.state
+
+    async def get_storage_data(self) -> SimpleNamespace:
+        fsm_state_data: Optional[FSMStateData] = await self._state_repo.get(
+            KEY_TEMPLATE.format(
+                host=self._message.bot.host,
+                bot_id=self._message.bot.id,
+                chat_id=self._message.chat.id,
+                user_huid=self._message.sender.huid,
+            ),
+        )
+
+        if not fsm_state_data:
+            return SimpleNamespace()
+
+        return fsm_state_data.storage
 
     async def drop_state(self) -> None:
         await self._state_repo.delete(

pyproject.toml

@@ -2,7 +2,7 @@
 
 [tool.poetry]
 name = "pybotx-fsm"
-version = "0.4.5"
+version = "0.5.0"
 description = "FSM middleware for using with pybotx"
 readme = "README.md"
 authors = []

tests/test_pass_args.py

@@ -80,3 +80,130 @@ async def fsm_handler(message: IncomingMessage, _: Bot) -> None:
     assert second_message.state.fsm_storage.first_message == "Hello Friend!"
     assert second_message.state.fsm_storage.second_message == "Hello World!"
     assert default_handler_trigger.called
+
+
+async def test_reset_args(
+    bot_account: BotAccountWithSecret,
+    incoming_message_factory: Callable[..., IncomingMessage],
+) -> None:
+    # - Arrange -
+    first_handler_trigger = Mock()
+    second_handler_trigger = Mock()
+    default_handler_trigger = Mock()
+
+    class EnumForTests(Enum):
+        FIRST_STATE = auto()
+        SECOND_STATE = auto()
+
+    collector = HandlerCollector()
+    fsm = FSMCollector(EnumForTests)
+
+    @fsm.on(EnumForTests.FIRST_STATE)
+    async def process_first_state(message: IncomingMessage, _: Bot) -> None:
+        await message.state.fsm.change_state(
+            EnumForTests.SECOND_STATE,
+            message="Hello World!",
+        )
+        first_handler_trigger()
+
+    @fsm.on(EnumForTests.SECOND_STATE)
+    async def process_second_state(message: IncomingMessage, _: Bot) -> None:
+        await message.state.fsm.drop_state()
+        second_handler_trigger()
+
+    @collector.default_message_handler
+    async def default_handler(__: IncomingMessage, _: Bot) -> None:
+        default_handler_trigger()
+
+    @collector.command("/fsm", visible=False)
+    async def fsm_handler(message: IncomingMessage, _: Bot) -> None:
+        await message.state.fsm.change_state(
+            EnumForTests.FIRST_STATE, message="Hello Friend!"
+        )
+
+    built_bot = Bot(
+        collectors=[collector],
+        bot_accounts=[bot_account],
+        middlewares=[FSMMiddleware([fsm], state_repo_key="state_repo")],
+    )
+    built_bot.state.state_repo = StateRepo()
+
+    init_message = incoming_message_factory(body="/fsm")
+    first_message = incoming_message_factory(body="text")
+    second_message = incoming_message_factory(body="text")
+    third_message = incoming_message_factory(body="text")
+
+    # - Act -
+    async with lifespan_wrapper(built_bot) as bot:
+        await bot.async_execute_bot_command(init_message)
+        await bot.async_execute_bot_command(first_message)
+        await bot.async_execute_bot_command(second_message)
+        await bot.async_execute_bot_command(third_message)
+
+    # - Assert -
+    assert first_message.state.fsm_storage.message == "Hello World!"
+    assert second_message.state.fsm_storage.message == "Hello World!"
+
+
+async def test_augment_args(
+    bot_account: BotAccountWithSecret,
+    incoming_message_factory: Callable[..., IncomingMessage],
+) -> None:
+    # - Arrange -
+    first_handler_trigger = Mock()
+    second_handler_trigger = Mock()
+    default_handler_trigger = Mock()
+
+    class EnumForTests(Enum):
+        FIRST_STATE = auto()
+        SECOND_STATE = auto()
+
+    collector = HandlerCollector()
+    fsm = FSMCollector(EnumForTests)
+
+    @fsm.on(EnumForTests.FIRST_STATE)
+    async def process_first_state(message: IncomingMessage, _: Bot) -> None:
+        await message.state.fsm.change_state(
+            EnumForTests.SECOND_STATE,
+            second_message="Hello World!",
+        )
+        first_handler_trigger()
+
+    @fsm.on(EnumForTests.SECOND_STATE)
+    async def process_second_state(message: IncomingMessage, _: Bot) -> None:
+        await message.state.fsm.drop_state()
+        second_handler_trigger()
+
+    @collector.default_message_handler
+    async def default_handler(__: IncomingMessage, _: Bot) -> None:
+        default_handler_trigger()
+
+    @collector.command("/fsm", visible=False)
+    async def fsm_handler(message: IncomingMessage, _: Bot) -> None:
+        await message.state.fsm.change_state(
+            EnumForTests.FIRST_STATE, first_message="Hello Friend!"
+        )
+
+    built_bot = Bot(
+        collectors=[collector],
+        bot_accounts=[bot_account],
+        middlewares=[FSMMiddleware([fsm], state_repo_key="state_repo")],
+    )
+    built_bot.state.state_repo = StateRepo()
+
+    init_message = incoming_message_factory(body="/fsm")
+    first_message = incoming_message_factory(body="text")
+    second_message = incoming_message_factory(body="text")
+    third_message = incoming_message_factory(body="text")
+
+    # - Act -
+    async with lifespan_wrapper(built_bot) as bot:
+        await bot.async_execute_bot_command(init_message)
+        await bot.async_execute_bot_command(first_message)
+        await bot.async_execute_bot_command(second_message)
+        await bot.async_execute_bot_command(third_message)
+
+    # - Assert -
+    assert first_message.state.fsm_storage.first_message == "Hello Friend!"
+    assert second_message.state.fsm_storage.first_message == "Hello Friend!"
+    assert second_message.state.fsm_storage.second_message == "Hello World!"