mirror of
https://github.com/python-telegram-bot/python-telegram-bot.git
synced 2024-12-28 23:38:31 +01:00
Add Method drop_chat/user_data to Dispatcher and Persistence (#2852)
This commit is contained in:
parent
e442782d8f
commit
0ccd7d40ac
7 changed files with 336 additions and 118 deletions
|
@ -457,7 +457,7 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
|
|||
|
||||
Returns:
|
||||
Optional[Tuple[List[Tuple[:obj:`str`, :obj:`float`, \
|
||||
Dict[:obj:`str`, :obj:`Any`]]], Dict[:obj:`str`, :obj:`str`]]:
|
||||
Dict[:obj:`str`, :obj:`Any`]]], Dict[:obj:`str`, :obj:`str`]]]:
|
||||
The restored meta data or :obj:`None`, if no data was stored.
|
||||
"""
|
||||
|
||||
|
@ -520,6 +520,44 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
|
|||
The :attr:`telegram.ext.Dispatcher.bot_data`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def update_callback_data(self, data: CDCData) -> None:
|
||||
"""Will be called by the :class:`telegram.ext.Dispatcher` after a handler has
|
||||
handled an update.
|
||||
|
||||
.. versionadded:: 13.6
|
||||
|
||||
.. versionchanged:: 14.0
|
||||
Changed this method into an ``@abstractmethod``.
|
||||
|
||||
Args:
|
||||
data (Optional[Tuple[List[Tuple[:obj:`str`, :obj:`float`, \
|
||||
Dict[:obj:`str`, :obj:`Any`]]], Dict[:obj:`str`, :obj:`str`]]]):
|
||||
The relevant data to restore :class:`telegram.ext.CallbackDataCache`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def drop_chat_data(self, chat_id: int) -> None:
|
||||
"""Will be called by the :class:`telegram.ext.Dispatcher`, when using
|
||||
:meth:`~telegram.ext.Dispatcher.drop_chat_data`.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
chat_id (:obj:`int`): The chat id to delete from the persistence.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def drop_user_data(self, user_id: int) -> None:
|
||||
"""Will be called by the :class:`telegram.ext.Dispatcher`, when using
|
||||
:meth:`~telegram.ext.Dispatcher.drop_user_data`.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
user_id (:obj:`int`): The user id to delete from the persistence.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def refresh_user_data(self, user_id: int, user_data: UD) -> None:
|
||||
"""Will be called by the :class:`telegram.ext.Dispatcher` before passing the
|
||||
|
@ -570,22 +608,6 @@ class BasePersistence(Generic[UD, CD, BD], ABC):
|
|||
The ``bot_data``.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def update_callback_data(self, data: CDCData) -> None:
|
||||
"""Will be called by the :class:`telegram.ext.Dispatcher` after a handler has
|
||||
handled an update.
|
||||
|
||||
.. versionadded:: 13.6
|
||||
|
||||
.. versionchanged:: 14.0
|
||||
Changed this method into an ``@abstractmethod``.
|
||||
|
||||
Args:
|
||||
data (Optional[Tuple[List[Tuple[:obj:`str`, :obj:`float`, \
|
||||
Dict[:obj:`str`, :obj:`Any`]]], Dict[:obj:`str`, :obj:`str`]]):
|
||||
The relevant data to restore :class:`telegram.ext.CallbackDataCache`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def flush(self) -> None:
|
||||
"""Will be called by :class:`telegram.ext.Updater` upon receiving a stop signal. Gives the
|
||||
|
|
|
@ -361,6 +361,32 @@ class DictPersistence(BasePersistence):
|
|||
self._callback_data = (data[0], data[1].copy())
|
||||
self._callback_data_json = None
|
||||
|
||||
def drop_chat_data(self, chat_id: int) -> None:
|
||||
"""Will delete the specified key from the :attr:`chat_data`.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
chat_id (:obj:`int`): The chat id to delete from the persistence.
|
||||
"""
|
||||
if self._chat_data is None:
|
||||
return
|
||||
self._chat_data.pop(chat_id, None)
|
||||
self._chat_data_json = None
|
||||
|
||||
def drop_user_data(self, user_id: int) -> None:
|
||||
"""Will delete the specified key from the :attr:`user_data`.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
user_id (:obj:`int`): The user id to delete from the persistence.
|
||||
"""
|
||||
if self._user_data is None:
|
||||
return
|
||||
self._user_data.pop(user_id, None)
|
||||
self._user_data_json = None
|
||||
|
||||
def refresh_user_data(self, user_id: int, user_data: Dict) -> None:
|
||||
"""Does nothing.
|
||||
|
||||
|
|
|
@ -37,7 +37,9 @@ from typing import (
|
|||
TypeVar,
|
||||
TYPE_CHECKING,
|
||||
Tuple,
|
||||
Mapping,
|
||||
)
|
||||
from types import MappingProxyType
|
||||
from uuid import uuid4
|
||||
|
||||
from telegram import Update
|
||||
|
@ -78,11 +80,11 @@ class DispatcherHandlerStop(Exception):
|
|||
Note:
|
||||
Has no effect, if the handler or error handler is run asynchronously.
|
||||
|
||||
Attributes:
|
||||
state (:obj:`object`): Optional. The next state of the conversation.
|
||||
|
||||
Args:
|
||||
state (:obj:`object`, optional): The next state of the conversation.
|
||||
|
||||
Attributes:
|
||||
state (:obj:`object`): Optional. The next state of the conversation.
|
||||
"""
|
||||
|
||||
__slots__ = ('state',)
|
||||
|
@ -111,8 +113,24 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
instance to pass onto handler callbacks.
|
||||
workers (:obj:`int`, optional): Number of maximum concurrent worker threads for the
|
||||
``@run_async`` decorator and :meth:`run_async`.
|
||||
user_data (:obj:`defaultdict`): A dictionary handlers can use to store data for the user.
|
||||
chat_data (:obj:`defaultdict`): A dictionary handlers can use to store data for the chat.
|
||||
chat_data (:obj:`types.MappingProxyType`): A dictionary handlers can use to store data for
|
||||
the chat.
|
||||
|
||||
.. versionchanged:: 14.0
|
||||
:attr:`chat_data` is now read-only
|
||||
|
||||
.. tip::
|
||||
Manually modifying :attr:`chat_data` is almost never needed and unadvisable.
|
||||
|
||||
user_data (:obj:`types.MappingProxyType`): A dictionary handlers can use to store data for
|
||||
the user.
|
||||
|
||||
.. versionchanged:: 14.0
|
||||
:attr:`user_data` is now read-only
|
||||
|
||||
.. tip::
|
||||
Manually modifying :attr:`user_data` is almost never needed and unadvisable.
|
||||
|
||||
bot_data (:obj:`dict`): A dictionary handlers can use to store data for the bot.
|
||||
persistence (:class:`telegram.ext.BasePersistence`): Optional. The persistence class to
|
||||
store data that should be persistent over restarts.
|
||||
|
@ -144,7 +162,9 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
'persistence',
|
||||
'update_queue',
|
||||
'job_queue',
|
||||
'_user_data',
|
||||
'user_data',
|
||||
'_chat_data',
|
||||
'chat_data',
|
||||
'bot_data',
|
||||
'_update_persistence_lock',
|
||||
|
@ -198,10 +218,15 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
stacklevel=stack_level,
|
||||
)
|
||||
|
||||
self.user_data: DefaultDict[int, UD] = defaultdict(self.context_types.user_data)
|
||||
self.chat_data: DefaultDict[int, CD] = defaultdict(self.context_types.chat_data)
|
||||
self._user_data: DefaultDict[int, UD] = defaultdict(self.context_types.user_data)
|
||||
self._chat_data: DefaultDict[int, CD] = defaultdict(self.context_types.chat_data)
|
||||
# Read only mapping-
|
||||
self.user_data: Mapping[int, UD] = MappingProxyType(self._user_data)
|
||||
self.chat_data: Mapping[int, CD] = MappingProxyType(self._chat_data)
|
||||
|
||||
self.bot_data = self.context_types.bot_data()
|
||||
self.persistence: Optional[BasePersistence] = None
|
||||
|
||||
self.persistence: Optional[BasePersistence]
|
||||
self._update_persistence_lock = Lock()
|
||||
if persistence:
|
||||
if not isinstance(persistence, BasePersistence):
|
||||
|
@ -213,13 +238,9 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
self.persistence.set_bot(self.bot)
|
||||
|
||||
if self.persistence.store_data.user_data:
|
||||
self.user_data = self.persistence.get_user_data()
|
||||
if not isinstance(self.user_data, defaultdict):
|
||||
raise ValueError("user_data must be of type defaultdict")
|
||||
self._user_data.update(self.persistence.get_user_data())
|
||||
if self.persistence.store_data.chat_data:
|
||||
self.chat_data = self.persistence.get_chat_data()
|
||||
if not isinstance(self.chat_data, defaultdict):
|
||||
raise ValueError("chat_data must be of type defaultdict")
|
||||
self._chat_data.update(self.persistence.get_chat_data())
|
||||
if self.persistence.store_data.bot_data:
|
||||
self.bot_data = self.persistence.get_bot_data()
|
||||
if not isinstance(self.bot_data, self.context_types.bot_data):
|
||||
|
@ -230,7 +251,7 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
persistent_data = self.persistence.get_callback_data()
|
||||
if persistent_data is not None:
|
||||
if not isinstance(persistent_data, tuple) and len(persistent_data) != 2:
|
||||
raise ValueError('callback_data must be a 2-tuple')
|
||||
raise ValueError('callback_data must be a tuple of length 2')
|
||||
# Mypy doesn't know that persistence.set_bot (see above) already checks that
|
||||
# self.bot is an instance of ExtBot if callback_data should be stored ...
|
||||
self.bot.callback_data_cache = CallbackDataCache( # type: ignore[attr-defined]
|
||||
|
@ -631,6 +652,34 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
if not self.handlers[group]:
|
||||
del self.handlers[group]
|
||||
|
||||
def drop_chat_data(self, chat_id: int) -> None:
|
||||
"""Used for deleting a key from the :attr:`chat_data`.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
chat_id (:obj:`int`): The chat id to delete from the persistence. The entry
|
||||
will be deleted even if it is not empty.
|
||||
"""
|
||||
self._chat_data.pop(chat_id, None) # type: ignore[arg-type]
|
||||
|
||||
if self.persistence:
|
||||
self.persistence.drop_chat_data(chat_id)
|
||||
|
||||
def drop_user_data(self, user_id: int) -> None:
|
||||
"""Used for deleting a key from the :attr:`user_data`.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
user_id (:obj:`int`): The user id to delete from the persistence. The entry
|
||||
will be deleted even if it is not empty.
|
||||
"""
|
||||
self._user_data.pop(user_id, None) # type: ignore[arg-type]
|
||||
|
||||
if self.persistence:
|
||||
self.persistence.drop_user_data(user_id)
|
||||
|
||||
def update_persistence(self, update: object = None) -> None:
|
||||
"""Update :attr:`user_data`, :attr:`chat_data` and :attr:`bot_data` in :attr:`persistence`.
|
||||
|
||||
|
@ -643,7 +692,7 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
|
||||
def __update_persistence(self, update: object = None) -> None:
|
||||
if self.persistence:
|
||||
# We use list() here in order to decouple chat_ids from self.chat_data, as dict view
|
||||
# We use list() here in order to decouple chat_ids from self._chat_data, as dict view
|
||||
# objects will change, when the dict does and we want to loop over chat_ids
|
||||
chat_ids = list(self.chat_data.keys())
|
||||
user_ids = list(self.user_data.keys())
|
||||
|
@ -758,7 +807,7 @@ class Dispatcher(Generic[BT, CCT, UD, CD, BD, JQ, PT]):
|
|||
|
||||
Returns:
|
||||
:obj:`bool`: :obj:`True` if one of the error handlers raised
|
||||
:class:`telegram.ext.DispatcherHandlerStop`. :obj:`False`, otherwise.
|
||||
:class:`telegram.ext.DispatcherHandlerStop`. :obj:`False`, otherwise.
|
||||
"""
|
||||
async_args = None if not promise else promise.args
|
||||
async_kwargs = None if not promise else promise.kwargs
|
||||
|
|
|
@ -392,6 +392,44 @@ class PicklePersistence(BasePersistence[UD, CD, BD]):
|
|||
else:
|
||||
self._dump_singlefile()
|
||||
|
||||
def drop_chat_data(self, chat_id: int) -> None:
|
||||
"""Will delete the specified key from the :attr:`chat_data` and depending on
|
||||
:attr:`on_flush` save the pickle file.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
chat_id (:obj:`int`): The chat id to delete from the persistence.
|
||||
"""
|
||||
if self.chat_data is None:
|
||||
return
|
||||
self.chat_data.pop(chat_id, None) # type: ignore[arg-type]
|
||||
|
||||
if not self.on_flush:
|
||||
if not self.single_file:
|
||||
self._dump_file(Path(f"{self.filepath}_chat_data"), self.chat_data)
|
||||
else:
|
||||
self._dump_singlefile()
|
||||
|
||||
def drop_user_data(self, user_id: int) -> None:
|
||||
"""Will delete the specified key from the :attr:`user_data` and depending on
|
||||
:attr:`on_flush` save the pickle file.
|
||||
|
||||
.. versionadded:: 14.0
|
||||
|
||||
Args:
|
||||
user_id (:obj:`int`): The user id to delete from the persistence.
|
||||
"""
|
||||
if self.user_data is None:
|
||||
return
|
||||
self.user_data.pop(user_id, None) # type: ignore[arg-type]
|
||||
|
||||
if not self.on_flush:
|
||||
if not self.single_file:
|
||||
self._dump_file(Path(f"{self.filepath}_user_data"), self.user_data)
|
||||
else:
|
||||
self._dump_singlefile()
|
||||
|
||||
def refresh_user_data(self, user_id: int, user_data: UD) -> None:
|
||||
"""Does nothing.
|
||||
|
||||
|
|
|
@ -28,6 +28,7 @@ from queue import Queue
|
|||
from threading import Thread, Event
|
||||
from time import sleep
|
||||
from typing import Callable, List, Iterable, Any
|
||||
from types import MappingProxyType
|
||||
|
||||
import pytest
|
||||
import pytz
|
||||
|
@ -194,10 +195,11 @@ def dp(_dp):
|
|||
# Reset the dispatcher first
|
||||
while not _dp.update_queue.empty():
|
||||
_dp.update_queue.get(False)
|
||||
_dp.chat_data = defaultdict(dict)
|
||||
_dp.user_data = defaultdict(dict)
|
||||
_dp._chat_data = defaultdict(dict)
|
||||
_dp._user_data = defaultdict(dict)
|
||||
_dp.chat_data = MappingProxyType(_dp._chat_data) # Rebuild the mapping so it updates
|
||||
_dp.user_data = MappingProxyType(_dp._user_data)
|
||||
_dp.bot_data = {}
|
||||
_dp.persistence = None
|
||||
_dp.handlers = {}
|
||||
_dp.error_handlers = {}
|
||||
_dp.exception_event = Event()
|
||||
|
|
|
@ -125,6 +125,15 @@ class TestDispatcher:
|
|||
)
|
||||
assert recwarn[0].filename == __file__, "stacklevel is incorrect!"
|
||||
|
||||
@pytest.mark.parametrize("data", ["chat_data", "user_data"])
|
||||
def test_chat_user_data_read_only(self, dp, data):
|
||||
read_only_data = getattr(dp, data)
|
||||
writable_data = getattr(dp, f"_{data}")
|
||||
writable_data[123] = 321
|
||||
assert read_only_data == writable_data
|
||||
with pytest.raises(TypeError):
|
||||
read_only_data[111] = 123
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'builder',
|
||||
(DispatcherBuilder(), UpdaterBuilder()),
|
||||
|
@ -621,6 +630,12 @@ class TestDispatcher:
|
|||
def update_bot_data(self, data):
|
||||
raise Exception
|
||||
|
||||
def drop_chat_data(self, chat_id):
|
||||
pass
|
||||
|
||||
def drop_user_data(self, user_id):
|
||||
pass
|
||||
|
||||
def get_chat_data(self):
|
||||
return defaultdict(dict)
|
||||
|
||||
|
@ -747,6 +762,12 @@ class TestDispatcher:
|
|||
def update_user_data(self, user_id, data):
|
||||
self.update(data)
|
||||
|
||||
def drop_user_data(self, user_id):
|
||||
pass
|
||||
|
||||
def drop_chat_data(self, chat_id):
|
||||
pass
|
||||
|
||||
def get_chat_data(self):
|
||||
pass
|
||||
|
||||
|
@ -825,6 +846,12 @@ class TestDispatcher:
|
|||
def update_conversation(self, name, key, new_state):
|
||||
pass
|
||||
|
||||
def drop_chat_data(self, chat_id):
|
||||
pass
|
||||
|
||||
def drop_user_data(self, user_id):
|
||||
pass
|
||||
|
||||
def get_conversations(self, name):
|
||||
pass
|
||||
|
||||
|
@ -879,6 +906,26 @@ class TestDispatcher:
|
|||
assert not dp.persistence.test_flag_user_data
|
||||
assert dp.persistence.test_flag_chat_data
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"c_id,expected",
|
||||
[(321, {222: "remove_me"}), (111, {321: {'not_empty': 'no'}, 222: "remove_me"})],
|
||||
ids=["test chat_id removal", "test no key in data (no error)"],
|
||||
)
|
||||
def test_drop_chat_data(self, dp, c_id, expected):
|
||||
dp._chat_data.update({321: {'not_empty': 'no'}, 222: "remove_me"})
|
||||
dp.drop_chat_data(c_id)
|
||||
assert dp.chat_data == expected
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"u_id,expected",
|
||||
[(321, {222: "remove_me"}), (111, {321: {'not_empty': 'no'}, 222: "remove_me"})],
|
||||
ids=["test user_id removal", "test no key in data (no error)"],
|
||||
)
|
||||
def test_drop_user_data(self, dp, u_id, expected):
|
||||
dp._user_data.update({321: {'not_empty': 'no'}, 222: "remove_me"})
|
||||
dp.drop_user_data(u_id)
|
||||
assert dp.user_data == expected
|
||||
|
||||
def test_update_persistence_once_per_update(self, monkeypatch, dp):
|
||||
def update_persistence(*args, **kwargs):
|
||||
self.count += 1
|
||||
|
|
|
@ -50,14 +50,14 @@ from telegram.ext import (
|
|||
JobQueue,
|
||||
ContextTypes,
|
||||
)
|
||||
from telegram.ext._callbackdatacache import _KeyboardData
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def change_directory(tmp_path):
|
||||
def change_directory(tmp_path: Path):
|
||||
orig_dir = Path.cwd()
|
||||
# Switch to a temporary directory so we don't have to worry about cleaning up files
|
||||
# (str() for py<3.6)
|
||||
os.chdir(str(tmp_path))
|
||||
# Switch to a temporary directory, so we don't have to worry about cleaning up files
|
||||
os.chdir(tmp_path)
|
||||
yield
|
||||
# Go back to original directory
|
||||
os.chdir(orig_dir)
|
||||
|
@ -99,6 +99,12 @@ class OwnPersistence(BasePersistence):
|
|||
def get_callback_data(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def drop_user_data(self, user_id):
|
||||
raise NotImplementedError
|
||||
|
||||
def drop_chat_data(self, chat_id):
|
||||
raise NotImplementedError
|
||||
|
||||
def refresh_user_data(self, user_id, user_data):
|
||||
raise NotImplementedError
|
||||
|
||||
|
@ -159,6 +165,12 @@ def bot_persistence():
|
|||
def update_callback_data(self, data):
|
||||
self.callback_data = data
|
||||
|
||||
def drop_user_data(self, user_id):
|
||||
pass
|
||||
|
||||
def drop_chat_data(self, chat_id):
|
||||
pass
|
||||
|
||||
def update_conversation(self, name, key, new_state):
|
||||
raise NotImplementedError
|
||||
|
||||
|
@ -257,7 +269,7 @@ class TestBasePersistence:
|
|||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
'flush, get_bot_data, get_callback_data, '
|
||||
'drop_chat_data, drop_user_data, flush, get_bot_data, get_callback_data, '
|
||||
'get_chat_data, get_conversations, '
|
||||
'get_user_data, refresh_bot_data, refresh_chat_data, '
|
||||
'refresh_user_data, update_bot_data, update_callback_data, '
|
||||
|
@ -284,52 +296,40 @@ class TestBasePersistence:
|
|||
def test_dispatcher_integration_init(
|
||||
self, bot, base_persistence, chat_data, user_data, bot_data, callback_data
|
||||
):
|
||||
def get_user_data():
|
||||
# Bad data testing-
|
||||
def bad_get_bot_data():
|
||||
return "test"
|
||||
|
||||
def get_chat_data():
|
||||
def bad_get_callback_data():
|
||||
return "test"
|
||||
|
||||
def get_bot_data():
|
||||
return "test"
|
||||
|
||||
def get_callback_data():
|
||||
return "test"
|
||||
|
||||
base_persistence.get_user_data = get_user_data
|
||||
base_persistence.get_chat_data = get_chat_data
|
||||
base_persistence.get_bot_data = get_bot_data
|
||||
base_persistence.get_callback_data = get_callback_data
|
||||
|
||||
with pytest.raises(ValueError, match="user_data must be of type defaultdict"):
|
||||
UpdaterBuilder().bot(bot).persistence(base_persistence).build()
|
||||
|
||||
def get_user_data():
|
||||
# Good data testing-
|
||||
def good_get_user_data():
|
||||
return user_data
|
||||
|
||||
base_persistence.get_user_data = get_user_data
|
||||
with pytest.raises(ValueError, match="chat_data must be of type defaultdict"):
|
||||
UpdaterBuilder().bot(bot).persistence(base_persistence).build()
|
||||
|
||||
def get_chat_data():
|
||||
def good_get_chat_data():
|
||||
return chat_data
|
||||
|
||||
base_persistence.get_chat_data = get_chat_data
|
||||
def good_get_bot_data():
|
||||
return bot_data
|
||||
|
||||
def good_get_callback_data():
|
||||
return callback_data
|
||||
|
||||
base_persistence.get_user_data = good_get_user_data # No errors to be tested so
|
||||
base_persistence.get_chat_data = good_get_chat_data
|
||||
base_persistence.get_bot_data = bad_get_bot_data
|
||||
base_persistence.get_callback_data = bad_get_callback_data
|
||||
|
||||
with pytest.raises(ValueError, match="bot_data must be of type dict"):
|
||||
UpdaterBuilder().bot(bot).persistence(base_persistence).build()
|
||||
|
||||
def get_bot_data():
|
||||
return bot_data
|
||||
|
||||
base_persistence.get_bot_data = get_bot_data
|
||||
with pytest.raises(ValueError, match="callback_data must be a 2-tuple"):
|
||||
base_persistence.get_bot_data = good_get_bot_data
|
||||
with pytest.raises(ValueError, match="callback_data must be a tuple of length 2"):
|
||||
UpdaterBuilder().bot(bot).persistence(base_persistence).build()
|
||||
|
||||
def get_callback_data():
|
||||
return callback_data
|
||||
|
||||
base_persistence.bot = None
|
||||
base_persistence.get_callback_data = get_callback_data
|
||||
base_persistence.get_callback_data = good_get_callback_data
|
||||
u = UpdaterBuilder().bot(bot).persistence(base_persistence).build()
|
||||
assert u.dispatcher.bot is base_persistence.bot
|
||||
assert u.dispatcher.bot_data == bot_data
|
||||
|
@ -339,7 +339,7 @@ class TestBasePersistence:
|
|||
u.dispatcher.chat_data[442233]['test5'] = 'test6'
|
||||
assert u.dispatcher.chat_data[442233]['test5'] == 'test6'
|
||||
|
||||
@pytest.mark.parametrize('run_async', [True, False], ids=['synchronous', 'run_async'])
|
||||
@pytest.mark.parametrize('run_async', [True, False], ids=['run_async', 'synchronous'])
|
||||
def test_dispatcher_integration_handlers(
|
||||
self,
|
||||
dp,
|
||||
|
@ -368,8 +368,6 @@ class TestBasePersistence:
|
|||
base_persistence.get_chat_data = get_chat_data
|
||||
base_persistence.get_bot_data = get_bot_data
|
||||
base_persistence.get_callback_data = get_callback_data
|
||||
# base_persistence.update_chat_data = lambda x: x
|
||||
# base_persistence.update_user_data = lambda x: x
|
||||
base_persistence.refresh_bot_data = lambda x: x
|
||||
base_persistence.refresh_chat_data = lambda x, y: x
|
||||
base_persistence.refresh_user_data = lambda x, y: x
|
||||
|
@ -383,7 +381,7 @@ class TestBasePersistence:
|
|||
pytest.fail('bot_data corrupt')
|
||||
|
||||
def callback_known_chat(update, context):
|
||||
if not context.chat_data['test3'] == 'test4':
|
||||
if not context.chat_data[3] == 'test4':
|
||||
pytest.fail('chat_data corrupt')
|
||||
if not context.bot_data == bot_data:
|
||||
pytest.fail('bot_data corrupt')
|
||||
|
@ -398,20 +396,17 @@ class TestBasePersistence:
|
|||
context.user_data[1] = 'test7'
|
||||
context.chat_data[2] = 'test8'
|
||||
context.bot_data['test0'] = 'test0'
|
||||
context.bot.callback_data_cache.put('test0')
|
||||
# Let's now delete user1 and chat1
|
||||
context.dispatcher.drop_chat_data(-67890)
|
||||
context.dispatcher.drop_user_data(12345)
|
||||
# Test setting new keyboard callback data-
|
||||
context.bot.callback_data_cache._keyboard_data['id'] = _KeyboardData(
|
||||
'id', button_data={'button3': 'test3'}
|
||||
)
|
||||
|
||||
known_user = MessageHandler(
|
||||
filters.User(user_id=12345),
|
||||
callback_known_user,
|
||||
)
|
||||
known_chat = MessageHandler(
|
||||
filters.Chat(chat_id=-67890),
|
||||
callback_known_chat,
|
||||
)
|
||||
unknown = MessageHandler(
|
||||
filters.ALL,
|
||||
callback_unknown_user_or_chat,
|
||||
)
|
||||
known_user = MessageHandler(filters.User(user_id=12345), callback_known_user) # user1
|
||||
known_chat = MessageHandler(filters.Chat(chat_id=-67890), callback_known_chat) # chat1
|
||||
unknown = MessageHandler(filters.ALL, callback_unknown_user_or_chat) # user2 and chat2
|
||||
dp.add_handler(known_user)
|
||||
dp.add_handler(known_chat)
|
||||
dp.add_handler(unknown)
|
||||
|
@ -420,51 +415,64 @@ class TestBasePersistence:
|
|||
chat1 = Chat(id=-67890, type='group')
|
||||
chat2 = Chat(id=-987654, type='group')
|
||||
m = Message(1, None, chat2, from_user=user1)
|
||||
u = Update(0, m)
|
||||
with caplog.at_level(logging.ERROR):
|
||||
dp.process_update(u)
|
||||
rec = caplog.records[-1]
|
||||
assert rec.getMessage() == 'No error handlers are registered, logging exception.'
|
||||
assert rec.levelname == 'ERROR'
|
||||
rec = caplog.records[-2]
|
||||
assert rec.getMessage() == 'No error handlers are registered, logging exception.'
|
||||
assert rec.levelname == 'ERROR'
|
||||
rec = caplog.records[-3]
|
||||
assert rec.getMessage() == 'No error handlers are registered, logging exception.'
|
||||
assert rec.levelname == 'ERROR'
|
||||
u_known_user = Update(0, m)
|
||||
dp.process_update(u_known_user)
|
||||
# 4 errors which arise since update_*_data are raising NotImplementedError here.
|
||||
assert len(caplog.records) == 4
|
||||
m.from_user = user2
|
||||
m.chat = chat1
|
||||
u = Update(1, m)
|
||||
dp.process_update(u)
|
||||
u_known_chat = Update(1, m)
|
||||
dp.process_update(u_known_chat)
|
||||
m.chat = chat2
|
||||
u = Update(2, m)
|
||||
u_unknown_user_or_chat = Update(2, m)
|
||||
|
||||
def save_bot_data(data):
|
||||
if 'test0' not in data:
|
||||
pytest.fail()
|
||||
|
||||
def save_chat_data(data):
|
||||
if -987654 not in data:
|
||||
def save_chat_data(_id, data):
|
||||
if 2 not in data: # data should be: {2: 'test8'}
|
||||
pytest.fail()
|
||||
|
||||
def save_user_data(data):
|
||||
if 54321 not in data:
|
||||
def save_user_data(_id, data):
|
||||
if 1 not in data: # data should be: {1: 'test7'}
|
||||
pytest.fail()
|
||||
|
||||
def save_callback_data(data):
|
||||
if not assert_data_in_cache(dp.bot.callback_data, 'test0'):
|
||||
if not assert_data_in_cache(dp.bot.callback_data_cache, 'test3'):
|
||||
pytest.fail()
|
||||
|
||||
# Functions to check deletion-
|
||||
def delete_user_data(user_id):
|
||||
if 12345 != user_id:
|
||||
pytest.fail("The id being deleted is not of user1's")
|
||||
user_data.pop(user_id, None)
|
||||
|
||||
def delete_chat_data(chat_id):
|
||||
if -67890 != chat_id:
|
||||
pytest.fail("The chat id being deleted is not of chat1's")
|
||||
chat_data.pop(chat_id, None)
|
||||
|
||||
base_persistence.update_chat_data = save_chat_data
|
||||
base_persistence.update_user_data = save_user_data
|
||||
base_persistence.update_bot_data = save_bot_data
|
||||
base_persistence.update_callback_data = save_callback_data
|
||||
dp.process_update(u)
|
||||
base_persistence.drop_chat_data = delete_chat_data
|
||||
base_persistence.drop_user_data = delete_user_data
|
||||
dp.process_update(u_unknown_user_or_chat)
|
||||
|
||||
# Test callback_unknown_user_or_chat worked correctly-
|
||||
assert dp.user_data[54321][1] == 'test7'
|
||||
assert dp.chat_data[-987654][2] == 'test8'
|
||||
assert dp.bot_data['test0'] == 'test0'
|
||||
assert assert_data_in_cache(dp.bot.callback_data_cache, 'test0')
|
||||
assert assert_data_in_cache(dp.bot.callback_data_cache, 'test3')
|
||||
assert 12345 not in dp.user_data # Tests if dp.drop_user_data worked or not
|
||||
assert -67890 not in dp.chat_data
|
||||
assert len(caplog.records) == 8 # Errors double since new update is processed.
|
||||
for r in caplog.records:
|
||||
assert issubclass(r.exc_info[0], NotImplementedError)
|
||||
assert r.getMessage() == 'No error handlers are registered, logging exception.'
|
||||
assert r.levelname == 'ERROR'
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'store_user_data', [True, False], ids=['store_user_data-True', 'store_user_data-False']
|
||||
|
@ -475,7 +483,7 @@ class TestBasePersistence:
|
|||
@pytest.mark.parametrize(
|
||||
'store_bot_data', [True, False], ids=['store_bot_data-True', 'store_bot_data-False']
|
||||
)
|
||||
@pytest.mark.parametrize('run_async', [True, False], ids=['synchronous', 'run_async'])
|
||||
@pytest.mark.parametrize('run_async', [True, False], ids=['run_async', 'synchronous'])
|
||||
def test_persistence_dispatcher_integration_refresh_data(
|
||||
self,
|
||||
dp,
|
||||
|
@ -1042,15 +1050,11 @@ class TestPicklePersistence:
|
|||
assert retrieved == bot_data
|
||||
|
||||
def test_no_files_present_multi_file(self, pickle_persistence):
|
||||
assert pickle_persistence.get_user_data() == defaultdict(dict)
|
||||
assert pickle_persistence.get_user_data() == defaultdict(dict)
|
||||
assert pickle_persistence.get_chat_data() == defaultdict(dict)
|
||||
assert pickle_persistence.get_chat_data() == defaultdict(dict)
|
||||
assert pickle_persistence.get_bot_data() == {}
|
||||
assert pickle_persistence.get_bot_data() == {}
|
||||
assert pickle_persistence.get_callback_data() is None
|
||||
assert pickle_persistence.get_conversations('noname') == {}
|
||||
assert pickle_persistence.get_conversations('noname') == {}
|
||||
|
||||
def test_no_files_present_single_file(self, pickle_persistence):
|
||||
pickle_persistence.single_file = True
|
||||
|
@ -1342,6 +1346,8 @@ class TestPicklePersistence:
|
|||
with Path('pickletest_user_data').open('rb') as f:
|
||||
user_data_test = defaultdict(dict, pickle.load(f))
|
||||
assert user_data_test == user_data
|
||||
pickle_persistence.drop_user_data(67890)
|
||||
assert 67890 not in pickle_persistence.get_user_data()
|
||||
|
||||
chat_data = pickle_persistence.get_chat_data()
|
||||
chat_data[-12345]['test3']['test4'] = 'test6'
|
||||
|
@ -1354,6 +1360,8 @@ class TestPicklePersistence:
|
|||
with Path('pickletest_chat_data').open('rb') as f:
|
||||
chat_data_test = defaultdict(dict, pickle.load(f))
|
||||
assert chat_data_test == chat_data
|
||||
pickle_persistence.drop_chat_data(-67890)
|
||||
assert -67890 not in pickle_persistence.get_chat_data()
|
||||
|
||||
bot_data = pickle_persistence.get_bot_data()
|
||||
bot_data['test3']['test4'] = 'test6'
|
||||
|
@ -1408,6 +1416,8 @@ class TestPicklePersistence:
|
|||
with Path('pickletest').open('rb') as f:
|
||||
user_data_test = defaultdict(dict, pickle.load(f)['user_data'])
|
||||
assert user_data_test == user_data
|
||||
pickle_persistence.drop_user_data(67890)
|
||||
assert 67890 not in pickle_persistence.get_user_data()
|
||||
|
||||
chat_data = pickle_persistence.get_chat_data()
|
||||
chat_data[-12345]['test3']['test4'] = 'test6'
|
||||
|
@ -1420,6 +1430,8 @@ class TestPicklePersistence:
|
|||
with Path('pickletest').open('rb') as f:
|
||||
chat_data_test = defaultdict(dict, pickle.load(f)['chat_data'])
|
||||
assert chat_data_test == chat_data
|
||||
pickle_persistence.drop_chat_data(-67890)
|
||||
assert -67890 not in pickle_persistence.get_chat_data()
|
||||
|
||||
bot_data = pickle_persistence.get_bot_data()
|
||||
bot_data['test3']['test4'] = 'test6'
|
||||
|
@ -1487,6 +1499,9 @@ class TestPicklePersistence:
|
|||
pickle_persistence.update_user_data(54321, user_data[54321])
|
||||
assert pickle_persistence.user_data == user_data
|
||||
|
||||
pickle_persistence.drop_user_data(0)
|
||||
assert pickle_persistence.user_data == user_data
|
||||
|
||||
with Path('pickletest_user_data').open('rb') as f:
|
||||
user_data_test = defaultdict(dict, pickle.load(f))
|
||||
assert not user_data_test == user_data
|
||||
|
@ -1498,6 +1513,9 @@ class TestPicklePersistence:
|
|||
pickle_persistence.update_chat_data(54321, chat_data[54321])
|
||||
assert pickle_persistence.chat_data == chat_data
|
||||
|
||||
pickle_persistence.drop_chat_data(0)
|
||||
assert pickle_persistence.user_data == user_data
|
||||
|
||||
with Path('pickletest_chat_data').open('rb') as f:
|
||||
chat_data_test = defaultdict(dict, pickle.load(f))
|
||||
assert not chat_data_test == chat_data
|
||||
|
@ -1905,6 +1923,12 @@ class TestPicklePersistence:
|
|||
assert isinstance(persistence.get_bot_data(), bd)
|
||||
assert persistence.get_bot_data() == 0
|
||||
|
||||
persistence.user_data = None
|
||||
persistence.chat_data = None
|
||||
persistence.drop_user_data(123)
|
||||
persistence.drop_chat_data(123)
|
||||
assert isinstance(persistence.get_user_data(), defaultdict)
|
||||
assert isinstance(persistence.get_chat_data(), defaultdict)
|
||||
persistence.user_data = None
|
||||
persistence.chat_data = None
|
||||
persistence.update_user_data(1, ud(1))
|
||||
|
@ -2132,6 +2156,11 @@ class TestDictPersistence:
|
|||
dict_persistence.update_user_data(12345, user_data[12345])
|
||||
assert dict_persistence.user_data == user_data
|
||||
assert dict_persistence.user_data_json == json.dumps(user_data)
|
||||
dict_persistence.drop_user_data(67890)
|
||||
assert 67890 not in dict_persistence.user_data
|
||||
dict_persistence._user_data = None
|
||||
dict_persistence.drop_user_data(123)
|
||||
assert isinstance(dict_persistence.get_user_data(), defaultdict)
|
||||
|
||||
chat_data = dict_persistence.get_chat_data()
|
||||
chat_data[-12345]['test3']['test4'] = 'test6'
|
||||
|
@ -2144,6 +2173,11 @@ class TestDictPersistence:
|
|||
dict_persistence.update_chat_data(-12345, chat_data[-12345])
|
||||
assert dict_persistence.chat_data == chat_data
|
||||
assert dict_persistence.chat_data_json == json.dumps(chat_data)
|
||||
dict_persistence.drop_chat_data(-67890)
|
||||
assert -67890 not in dict_persistence.chat_data
|
||||
dict_persistence._chat_data = None
|
||||
dict_persistence.drop_chat_data(123)
|
||||
assert isinstance(dict_persistence.get_chat_data(), defaultdict)
|
||||
|
||||
bot_data = dict_persistence.get_bot_data()
|
||||
bot_data['test3']['test4'] = 'test6'
|
||||
|
|
Loading…
Reference in a new issue