File: test_baseupdateprocessor.py

package info (click to toggle)
python-telegram-bot 22.3-1
  • links: PTS
  • area: main
  • in suites: sid
  • size: 11,060 kB
  • sloc: python: 90,298; makefile: 176; sh: 4
file content (196 lines) | stat: -rw-r--r-- 7,110 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
#!/usr/bin/env python
#
# A library that provides a Python interface to the Telegram Bot API
# Copyright (C) 2015-2025
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser Public License for more details.
#
# You should have received a copy of the GNU Lesser Public License
# along with this program.  If not, see [http://www.gnu.org/licenses/].
"""Here we run tests directly with SimpleUpdateProcessor because that's easier than providing dummy
implementations for SimpleUpdateProcessor and we want to test SimpleUpdateProcessor anyway."""
import asyncio

import pytest

from telegram import Update
from telegram.ext import SimpleUpdateProcessor
from tests.auxil.asyncio_helpers import call_after
from tests.auxil.slots import mro_slots


@pytest.fixture
def mock_processor():
    class MockProcessor(SimpleUpdateProcessor):
        test_flag = False

        async def do_process_update(self, update, coroutine):
            await coroutine
            self.test_flag = True

    return MockProcessor(5)


class TestSimpleUpdateProcessor:
    def test_slot_behaviour(self):
        inst = SimpleUpdateProcessor(1)
        for attr in inst.__slots__:
            assert getattr(inst, attr, "err") != "err", f"got extra slot '{attr}'"
        assert len(mro_slots(inst)) == len(set(mro_slots(inst))), "duplicate slot"

    @pytest.mark.parametrize("concurrent_updates", [0, -1])
    def test_init(self, concurrent_updates):
        processor = SimpleUpdateProcessor(3)
        assert processor.max_concurrent_updates == 3
        with pytest.raises(ValueError, match="must be a positive integer"):
            SimpleUpdateProcessor(concurrent_updates)

    async def test_process_update(self, mock_processor):
        """Test that process_update calls do_process_update."""
        update = Update(1)

        async def coroutine():
            pass

        await mock_processor.process_update(update, coroutine())
        # This flag is set in the mock processor in do_process_update, telling us that
        # do_process_update was called.
        assert mock_processor.test_flag

    async def test_do_process_update(self):
        """Test that do_process_update calls the coroutine."""
        processor = SimpleUpdateProcessor(1)
        update = Update(1)
        test_flag = False

        async def coroutine():
            nonlocal test_flag
            test_flag = True

        await processor.do_process_update(update, coroutine())
        assert test_flag

    async def test_max_concurrent_updates_enforcement(self, mock_processor):
        """Test that max_concurrent_updates is enforced, i.e. that the processor will run
        at most max_concurrent_updates coroutines at the same time."""
        count = 2 * mock_processor.max_concurrent_updates
        events = {i: asyncio.Event() for i in range(count)}
        queue = asyncio.Queue()
        for event in events.values():
            await queue.put(event)

        async def callback():
            await asyncio.sleep(0.5)
            (await queue.get()).set()

        # We start several calls to `process_update` at the same time, each of them taking
        # 0.5 seconds to complete. We know that they are completed when the corresponding
        # event is set.
        tasks = [
            asyncio.create_task(mock_processor.process_update(update=_, coroutine=callback()))
            for _ in range(count)
        ]

        # Right now we expect no event to be set
        for i in range(count):
            assert not events[i].is_set()

        # After 0.5 seconds (+ some buffer), we expect that exactly max_concurrent_updates
        # events are set.
        await asyncio.sleep(0.75)
        for i in range(mock_processor.max_concurrent_updates):
            assert events[i].is_set()
        for i in range(
            mock_processor.max_concurrent_updates,
            count,
        ):
            assert not events[i].is_set()

        # After wating another 0.5 seconds, we expect that the next max_concurrent_updates
        # events are set.
        await asyncio.sleep(0.5)
        for i in range(count):
            assert events[i].is_set()

        # Sanity check: we expect that all tasks are completed.
        await asyncio.gather(*tasks)

    async def test_context_manager(self, monkeypatch, mock_processor):
        self.test_flag = set()

        async def after_initialize(*args, **kwargs):
            self.test_flag.add("initialize")

        async def after_shutdown(*args, **kwargs):
            self.test_flag.add("stop")

        monkeypatch.setattr(
            SimpleUpdateProcessor,
            "initialize",
            call_after(SimpleUpdateProcessor.initialize, after_initialize),
        )
        monkeypatch.setattr(
            SimpleUpdateProcessor,
            "shutdown",
            call_after(SimpleUpdateProcessor.shutdown, after_shutdown),
        )

        async with mock_processor:
            pass

        assert self.test_flag == {"initialize", "stop"}

    async def test_context_manager_exception_on_init(self, monkeypatch, mock_processor):
        async def initialize(*args, **kwargs):
            raise RuntimeError("initialize")

        async def shutdown(*args, **kwargs):
            self.test_flag = "shutdown"

        monkeypatch.setattr(SimpleUpdateProcessor, "initialize", initialize)
        monkeypatch.setattr(SimpleUpdateProcessor, "shutdown", shutdown)

        with pytest.raises(RuntimeError, match="initialize"):
            async with mock_processor:
                pass

        assert self.test_flag == "shutdown"

    async def test_current_concurrent_updates(self, mock_processor):
        async def callback(event: asyncio.Event):
            await event.wait()

        events = {i: asyncio.Event() for i in range(10)}
        coroutines = {i: callback(event) for i, event in events.items()}

        process_tasks = [
            asyncio.create_task(mock_processor.process_update(Update(i), coroutines[i]))
            for i in range(10)
        ]
        await asyncio.sleep(0.01)

        assert mock_processor.current_concurrent_updates == mock_processor.max_concurrent_updates
        for i in range(5):
            events[i].set()

        await asyncio.sleep(0.01)
        assert mock_processor.current_concurrent_updates == mock_processor.max_concurrent_updates

        for i in range(5, 10):
            events[i].set()
            await asyncio.sleep(0.01)
            assert (
                mock_processor.current_concurrent_updates
                == mock_processor.max_concurrent_updates - (i - 4)
            )

        await asyncio.gather(*process_tasks)