Skip to content

Commit f7e81b9

Browse files
EldinnieEldinnie
committed
Continue work
Add tests for Basepersistence
1 parent 03056e1 commit f7e81b9

File tree

6 files changed

+252
-23
lines changed

6 files changed

+252
-23
lines changed

‎telegram/ext/basepersistence.py‎

Lines changed: 39 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,33 +20,63 @@
2020

2121

2222
class BasePersistence(object):
23-
def __init__(self, store_job_queue=False, store_user_data=False, store_chat_data=False):
24-
self.store_job_queue = store_job_queue
23+
def __init__(self, store_user_data=False, store_chat_data=False):
24+
"""
25+
Args:
26+
store_user_data (:obj:`bool`): Whether user_data should be saved by this
27+
persistence class.
28+
store_chat_data (:obj:`bool`): Whether user_data should be saved by this
29+
persistence class
30+
"""
2531
self.store_user_data = store_user_data
2632
self.store_chat_data = store_chat_data
2733

28-
def get_job_queue(self):
29-
raise NotImplementedError
30-
3134
def get_user_data(self):
35+
""""
36+
Returns:
37+
:obj:'defaultdict`: The restored user data.
38+
"""
3239
raise NotImplementedError
3340

3441
def get_chat_data(self):
42+
""""
43+
Returns:
44+
:obj:'defaultdict`: The restored chat data.
45+
"""
3546
raise NotImplementedError
3647

37-
def get_conversations(self):
38-
raise NotImplementedError
48+
def get_conversations(self, name):
49+
""""
50+
Args:
51+
name (:obj:`str`): The handlers name.
3952
40-
def update_conversation(self, conversations):
53+
Returns:
54+
:obj:'dict`: The restored conversations for the handler.
55+
"""
4156
raise NotImplementedError
4257

43-
def update_job_queue(self):
58+
def update_conversation(self, conversations):
59+
""""
60+
Args:
61+
conversation (:obj:'dict`): The :attr:`telegram.ext.ConversationHandler.conversations`
62+
dict to store.
63+
"""
4464
raise NotImplementedError
4565

4666
def update_user_data(self, user_data):
67+
""""
68+
Args:
69+
user_data (:obj:'defaultdict`): The :attr:`telegram.ext.dispatcher.user_data`
70+
defaultdict to store.
71+
"""
4772
raise NotImplementedError
4873

4974
def update_chat_data(self, chat_data):
75+
""""
76+
Args:
77+
chat_data (:obj:'defaultdict`): The :attr:`telegram.ext.dispatcher.chat_data`
78+
defaultdict to store.
79+
"""
5080
raise NotImplementedError
5181

5282
def flush(self):

‎telegram/ext/conversationhandler.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,10 @@ class ConversationHandler(Handler):
7676
per_user (:obj:`bool`): Optional. If the conversationkey should contain the User's ID.
7777
per_message (:obj:`bool`): Optional. If the conversationkey should contain the Message's
7878
ID.
79+
name (:obj:`str`): Optional. The name for this conversationhandler. Required for
80+
persistence
81+
persistent (:obj:`bool`): Optional. If the conversations dict for this handler should be
82+
saved. Name is required and persistence has to be set in :class:`Updater`
7983
8084
Args:
8185
entry_points (List[:class:`telegram.ext.Handler`]): A list of ``Handler`` objects that can
@@ -107,6 +111,10 @@ class ConversationHandler(Handler):
107111
Default is ``True``.
108112
per_message (:obj:`bool`, optional): If the conversationkey should contain the Message's
109113
ID. Default is ``False``.
114+
name (:obj:`str`, optional): The name for this conversationhandler. Required for
115+
persistence
116+
persistent (:obj:`bool`, optional): If the conversations dict for this handler should be
117+
saved. Name is required and persistence has to be set in :class:`Updater`
110118
111119
Raises:
112120
ValueError

‎telegram/ext/dispatcher.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,8 @@ class Dispatcher(object):
7474
decorator.
7575
user_data (:obj:`dict`): A dictionary handlers can use to store data for the user.
7676
chat_data (:obj:`dict`): A dictionary handlers can use to store data for the chat.
77+
persistence (:class:`telegram.ext.BasePersistence`): Optional. The persistence class to
78+
store data that should be persistent over restarts
7779
7880
Args:
7981
bot (:class:`telegram.Bot`): The bot object that should be passed to the handlers.
@@ -82,6 +84,8 @@ class Dispatcher(object):
8284
instance to pass onto handler callbacks.
8385
workers (:obj:`int`, optional): Number of maximum concurrent worker threads for the
8486
``@run_async`` decorator. defaults to 4.
87+
persistence (:class:`telegram.ext.BasePersistence`, optional): The persistence class to
88+
store data that should be persistent over restarts
8589
8690
"""
8791

@@ -301,9 +305,15 @@ def process_update(self, update):
301305
handler.handle_update(update, self)
302306
if self.persistence:
303307
if self.persistence.store_chat_data:
304-
self.persistence.update_chat_data(self.chat_data)
308+
try:
309+
self.persistence.update_chat_data(self.chat_data)
310+
except Exception:
311+
self.logger.exception('Saving chat data raised an error')
305312
if self.persistence.store_user_data:
306-
self.persistence.update_user_data(self.user_data)
313+
try:
314+
self.persistence.update_user_data(self.user_data)
315+
except Exception:
316+
self.logger.exception('Saving user data raised an error')
307317
break
308318

309319
# Stop processing with any other handler.
@@ -361,8 +371,8 @@ def add_handler(self, handler, group=DEFAULT_GROUP):
361371
if isinstance(handler, ConversationHandler) and handler.persistent:
362372
if not self.persistence:
363373
raise ValueError(
364-
"Conversationhandler {} can not be persistent if logger has no persistence",
365-
handler.name)
374+
"Conversationhandler {} can not be persistent if dispatcher has no "
375+
"persistence".format(handler.name))
366376
handler.conversations = self.persistence.get_conversations(handler.name)
367377
handler.persistence = self.persistence
368378

‎telegram/ext/updater.py‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@ class Updater(object):
5858
dispatcher (:class:`telegram.ext.Dispatcher`): Dispatcher that handles the updates and
5959
dispatches them to the handlers.
6060
running (:obj:`bool`): Indicates if the updater is running.
61+
persistence (:class:`telegram.ext.BasePersistence`): Optional. The persistence class to
62+
store data that should be persistent over restarts.
6163
6264
Args:
6365
token (:obj:`str`, optional): The bot's token given by the @BotFather.
@@ -74,6 +76,8 @@ class Updater(object):
7476
`telegram.utils.request.Request` object (ignored if `bot` argument is used). The
7577
request_kwargs are very useful for the advanced users who would like to control the
7678
default timeouts and/or control the proxy used for http communication.
79+
persistence (:class:`telegram.ext.BasePersistence`, optional): The persistence class to
80+
store data that should be persistent over restarts
7781
7882
Note:
7983
You must supply either a :attr:`bot` or a :attr:`token` argument.
@@ -124,17 +128,9 @@ def __init__(self,
124128
self.bot = Bot(token, base_url, request=self._request)
125129
self.user_sig_handler = user_sig_handler
126130
self.update_queue = Queue()
127-
self.persistence = persistence
128-
self.job_queue = None
129-
if self.persistence:
130-
if self.persistence.store_job_queue:
131-
self.job_queue = self.persistence.get_job_queue()
132-
if not isinstance(self.job_queue, JobQueue):
133-
raise ValueError("job_queue must be of type JobQueue")
134-
135-
if not self.job_queue:
136-
self.job_queue = JobQueue(self.bot)
131+
self.job_queue = JobQueue(self.bot)
137132
self.__exception_event = Event()
133+
self.persistence = persistence
138134
self.dispatcher = Dispatcher(
139135
self.bot,
140136
self.update_queue,

‎tests/test_conversationhandler.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,12 @@ def test_per_all_false(self):
9191
ConversationHandler(self.entry_points, self.states, self.fallbacks,
9292
per_chat=False, per_user=False, per_message=False)
9393

94+
def test_name_and_persistent(self, dp):
95+
with pytest.raises(ValueError, match="when handler is unnamed"):
96+
dp.add_handler(ConversationHandler([], {}, [], persistent=True))
97+
c = ConversationHandler([], {}, [], name="handler", persistent=True)
98+
assert c.name == "handler"
99+
94100
def test_conversation_handler(self, dp, bot, user1, user2):
95101
handler = ConversationHandler(entry_points=self.entry_points, states=self.states,
96102
fallbacks=self.fallbacks)

‎tests/test_persistence.py‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
#!/usr/bin/env python
2+
#
3+
# A library that provides a Python interface to the Telegram Bot API
4+
# Copyright (C) 2015-2018
5+
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
6+
#
7+
# This program is free software: you can redistribute it and/or modify
8+
# it under the terms of the GNU Lesser Public License as published by
9+
# the Free Software Foundation, either version 3 of the License, or
10+
# (at your option) any later version.
11+
#
12+
# This program is distributed in the hope that it will be useful,
13+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
14+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15+
# GNU Lesser Public License for more details.
16+
#
17+
# You should have received a copy of the GNU Lesser Public License
18+
# along with this program. If not, see [http://www.gnu.org/licenses/].
19+
import logging
20+
from collections import defaultdict
21+
22+
import pytest
23+
24+
from telegram import Update, Message, User, Chat
25+
from telegram.ext import BasePersistence, Updater, ConversationHandler, MessageHandler, Filters
26+
27+
28+
@pytest.fixture(scope="function")
29+
def base_persistence():
30+
return BasePersistence(store_chat_data=True, store_user_data=True)
31+
32+
33+
@pytest.fixture(scope="function")
34+
def chat_data():
35+
return defaultdict(dict, {-12345: {'test1': 'test2'}, -67890: {'test3': 'test4'}})
36+
37+
38+
@pytest.fixture(scope="function")
39+
def user_data():
40+
return defaultdict(dict, {12345: {'test1': 'test2'}, 67890: {'test3': 'test4'}})
41+
42+
43+
@pytest.fixture(scope="function")
44+
def updater(bot, base_persistence):
45+
base_persistence.store_chat_data = False
46+
base_persistence.store_user_data = False
47+
u = Updater(bot=bot, persistence=base_persistence)
48+
base_persistence.store_chat_data = True
49+
base_persistence.store_user_data = True
50+
return u
51+
52+
53+
class TestBasePersistence(object):
54+
55+
def test_creation(self, base_persistence):
56+
assert base_persistence.store_chat_data
57+
assert base_persistence.store_user_data
58+
with pytest.raises(NotImplementedError):
59+
base_persistence.get_chat_data()
60+
with pytest.raises(NotImplementedError):
61+
base_persistence.get_user_data()
62+
with pytest.raises(NotImplementedError):
63+
base_persistence.get_conversations("test")
64+
with pytest.raises(NotImplementedError):
65+
base_persistence.update_chat_data(None)
66+
with pytest.raises(NotImplementedError):
67+
base_persistence.update_user_data(None)
68+
with pytest.raises(NotImplementedError):
69+
base_persistence.update_conversation(None)
70+
with pytest.raises(NotImplementedError):
71+
base_persistence.flush()
72+
73+
def test_implementation(self, updater, base_persistence):
74+
dp = updater.dispatcher
75+
assert dp.persistence == base_persistence
76+
77+
def test_conversationhandler_addition(self, dp, base_persistence):
78+
with pytest.raises(ValueError, match="when handler is unnamed"):
79+
ConversationHandler([], [], [], persistent=True)
80+
with pytest.raises(ValueError, match="if dispatcher has no persistence"):
81+
dp.add_handler(ConversationHandler([], {}, [], persistent=True, name="My Handler"))
82+
dp.persistence = base_persistence
83+
with pytest.raises(NotImplementedError):
84+
dp.add_handler(ConversationHandler([], {}, [], persistent=True, name="My Handler"))
85+
86+
def test_dispatcher_integration_init(self, bot, base_persistence, chat_data, user_data):
87+
def get_user_data():
88+
return "test"
89+
90+
def get_chat_data():
91+
return "test"
92+
93+
base_persistence.get_user_data = get_user_data
94+
base_persistence.get_chat_data = get_chat_data
95+
with pytest.raises(ValueError, match="user_data must be of type defaultdict"):
96+
u = Updater(bot=bot, persistence=base_persistence)
97+
98+
def get_user_data():
99+
return user_data
100+
101+
base_persistence.get_user_data = get_user_data
102+
with pytest.raises(ValueError, match="chat_data must be of type defaultdict"):
103+
u = Updater(bot=bot, persistence=base_persistence)
104+
105+
def get_chat_data():
106+
return chat_data
107+
108+
base_persistence.get_chat_data = get_chat_data
109+
u = Updater(bot=bot, persistence=base_persistence)
110+
assert u.dispatcher.chat_data == chat_data
111+
assert u.dispatcher.user_data == user_data
112+
u.dispatcher.chat_data[442233]['test5'] = 'test6'
113+
assert u.dispatcher.chat_data[442233]['test5'] == 'test6'
114+
115+
def test_dispatcher_integration_handlers(self, caplog, bot, base_persistence: BasePersistence,
116+
chat_data, user_data):
117+
def get_user_data():
118+
return user_data
119+
120+
def get_chat_data():
121+
return chat_data
122+
123+
base_persistence.get_user_data = get_user_data
124+
base_persistence.get_chat_data = get_chat_data
125+
# base_persistence.update_chat_data = lambda x: x
126+
# base_persistence.update_user_data = lambda x: x
127+
updater = Updater(bot=bot, persistence=base_persistence)
128+
dp = updater.dispatcher
129+
130+
def callback_known_user(bot, update, user_data, chat_data):
131+
if not user_data['test1'] == 'test2':
132+
pytest.fail('user_data corrupt')
133+
134+
def callback_known_chat(bot, update, user_data, chat_data):
135+
if not chat_data['test3'] == 'test4':
136+
pytest.fail('chat_data corrupt')
137+
138+
def callback_unknown_user_or_chat(bot, update, user_data, chat_data):
139+
if not user_data == {}:
140+
pytest.fail('user_data corrupt')
141+
if not chat_data == {}:
142+
pytest.fail('chat_data corrupt')
143+
user_data[1] = 'test7'
144+
chat_data[2] = 'test8'
145+
146+
known_user = MessageHandler(Filters.user(user_id=12345), callback_known_user,
147+
pass_chat_data=True, pass_user_data=True)
148+
known_chat = MessageHandler(Filters.chat(chat_id=-67890), callback_known_chat,
149+
pass_chat_data=True, pass_user_data=True)
150+
unknown = MessageHandler(Filters.all, callback_unknown_user_or_chat, pass_chat_data=True,
151+
pass_user_data=True)
152+
dp.add_handler(known_user)
153+
dp.add_handler(known_chat)
154+
dp.add_handler(unknown)
155+
user1 = User(id=12345, first_name='test user', is_bot=False)
156+
user2 = User(id=54321, first_name='test user', is_bot=False)
157+
chat1 = Chat(id=-67890, type='group')
158+
chat2 = Chat(id=-987654, type='group')
159+
m = Message(1, user1, None, chat2)
160+
u = Update(0, m)
161+
with caplog.at_level(logging.ERROR):
162+
dp.process_update(u)
163+
rec = caplog.records[-1]
164+
assert rec.msg == 'Saving user data raised an error'
165+
assert rec.levelname == 'ERROR'
166+
rec = caplog.records[-2]
167+
assert rec.msg == 'Saving chat data raised an error'
168+
assert rec.levelname == 'ERROR'
169+
170+
m.from_user = user2
171+
m.chat = chat1
172+
u = Update(1, m)
173+
dp.process_update(u)
174+
m.chat = chat2
175+
u = Update(2, m)
176+
dp.process_update(u)
177+
178+
assert dp.user_data[54321][1] == 'test7'
179+
assert dp.chat_data[-987654][2] == 'test8'

0 commit comments

Comments
 (0)