• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import datetime
16import re
17
18import mock
19import pytest
20
21from google.api_core import exceptions
22from google.api_core import retry_async
23
24
25@mock.patch("asyncio.sleep", autospec=True)
26@mock.patch(
27    "google.api_core.datetime_helpers.utcnow",
28    return_value=datetime.datetime.min,
29    autospec=True,
30)
31@pytest.mark.asyncio
32async def test_retry_target_success(utcnow, sleep):
33    predicate = retry_async.if_exception_type(ValueError)
34    call_count = [0]
35
36    async def target():
37        call_count[0] += 1
38        if call_count[0] < 3:
39            raise ValueError()
40        return 42
41
42    result = await retry_async.retry_target(target, predicate, range(10), None)
43
44    assert result == 42
45    assert call_count[0] == 3
46    sleep.assert_has_calls([mock.call(0), mock.call(1)])
47
48
49@mock.patch("asyncio.sleep", autospec=True)
50@mock.patch(
51    "google.api_core.datetime_helpers.utcnow",
52    return_value=datetime.datetime.min,
53    autospec=True,
54)
55@pytest.mark.asyncio
56async def test_retry_target_w_on_error(utcnow, sleep):
57    predicate = retry_async.if_exception_type(ValueError)
58    call_count = {"target": 0}
59    to_raise = ValueError()
60
61    async def target():
62        call_count["target"] += 1
63        if call_count["target"] < 3:
64            raise to_raise
65        return 42
66
67    on_error = mock.Mock()
68
69    result = await retry_async.retry_target(
70        target, predicate, range(10), None, on_error=on_error
71    )
72
73    assert result == 42
74    assert call_count["target"] == 3
75
76    on_error.assert_has_calls([mock.call(to_raise), mock.call(to_raise)])
77    sleep.assert_has_calls([mock.call(0), mock.call(1)])
78
79
80@mock.patch("asyncio.sleep", autospec=True)
81@mock.patch(
82    "google.api_core.datetime_helpers.utcnow",
83    return_value=datetime.datetime.min,
84    autospec=True,
85)
86@pytest.mark.asyncio
87async def test_retry_target_non_retryable_error(utcnow, sleep):
88    predicate = retry_async.if_exception_type(ValueError)
89    exception = TypeError()
90    target = mock.Mock(side_effect=exception)
91
92    with pytest.raises(TypeError) as exc_info:
93        await retry_async.retry_target(target, predicate, range(10), None)
94
95    assert exc_info.value == exception
96    sleep.assert_not_called()
97
98
99@mock.patch("asyncio.sleep", autospec=True)
100@mock.patch("google.api_core.datetime_helpers.utcnow", autospec=True)
101@pytest.mark.asyncio
102async def test_retry_target_deadline_exceeded(utcnow, sleep):
103    predicate = retry_async.if_exception_type(ValueError)
104    exception = ValueError("meep")
105    target = mock.Mock(side_effect=exception)
106    # Setup the timeline so that the first call takes 5 seconds but the second
107    # call takes 6, which puts the retry over the deadline.
108    utcnow.side_effect = [
109        # The first call to utcnow establishes the start of the timeline.
110        datetime.datetime.min,
111        datetime.datetime.min + datetime.timedelta(seconds=5),
112        datetime.datetime.min + datetime.timedelta(seconds=11),
113    ]
114
115    with pytest.raises(exceptions.RetryError) as exc_info:
116        await retry_async.retry_target(target, predicate, range(10), deadline=10)
117
118    assert exc_info.value.cause == exception
119    assert exc_info.match("Deadline of 10.0s exceeded")
120    assert exc_info.match("last exception: meep")
121    assert target.call_count == 2
122
123
124@pytest.mark.asyncio
125async def test_retry_target_bad_sleep_generator():
126    with pytest.raises(ValueError, match="Sleep generator"):
127        await retry_async.retry_target(
128            mock.sentinel.target, mock.sentinel.predicate, [], None
129        )
130
131
132class TestAsyncRetry:
133    def test_constructor_defaults(self):
134        retry_ = retry_async.AsyncRetry()
135        assert retry_._predicate == retry_async.if_transient_error
136        assert retry_._initial == 1
137        assert retry_._maximum == 60
138        assert retry_._multiplier == 2
139        assert retry_._deadline == 120
140        assert retry_._on_error is None
141
142    def test_constructor_options(self):
143        _some_function = mock.Mock()
144
145        retry_ = retry_async.AsyncRetry(
146            predicate=mock.sentinel.predicate,
147            initial=1,
148            maximum=2,
149            multiplier=3,
150            deadline=4,
151            on_error=_some_function,
152        )
153        assert retry_._predicate == mock.sentinel.predicate
154        assert retry_._initial == 1
155        assert retry_._maximum == 2
156        assert retry_._multiplier == 3
157        assert retry_._deadline == 4
158        assert retry_._on_error is _some_function
159
160    def test_with_deadline(self):
161        retry_ = retry_async.AsyncRetry(
162            predicate=mock.sentinel.predicate,
163            initial=1,
164            maximum=2,
165            multiplier=3,
166            deadline=4,
167            on_error=mock.sentinel.on_error,
168        )
169        new_retry = retry_.with_deadline(42)
170        assert retry_ is not new_retry
171        assert new_retry._deadline == 42
172
173        # the rest of the attributes should remain the same
174        assert new_retry._predicate is retry_._predicate
175        assert new_retry._initial == retry_._initial
176        assert new_retry._maximum == retry_._maximum
177        assert new_retry._multiplier == retry_._multiplier
178        assert new_retry._on_error is retry_._on_error
179
180    def test_with_predicate(self):
181        retry_ = retry_async.AsyncRetry(
182            predicate=mock.sentinel.predicate,
183            initial=1,
184            maximum=2,
185            multiplier=3,
186            deadline=4,
187            on_error=mock.sentinel.on_error,
188        )
189        new_retry = retry_.with_predicate(mock.sentinel.predicate)
190        assert retry_ is not new_retry
191        assert new_retry._predicate == mock.sentinel.predicate
192
193        # the rest of the attributes should remain the same
194        assert new_retry._deadline == retry_._deadline
195        assert new_retry._initial == retry_._initial
196        assert new_retry._maximum == retry_._maximum
197        assert new_retry._multiplier == retry_._multiplier
198        assert new_retry._on_error is retry_._on_error
199
200    def test_with_delay_noop(self):
201        retry_ = retry_async.AsyncRetry(
202            predicate=mock.sentinel.predicate,
203            initial=1,
204            maximum=2,
205            multiplier=3,
206            deadline=4,
207            on_error=mock.sentinel.on_error,
208        )
209        new_retry = retry_.with_delay()
210        assert retry_ is not new_retry
211        assert new_retry._initial == retry_._initial
212        assert new_retry._maximum == retry_._maximum
213        assert new_retry._multiplier == retry_._multiplier
214
215    def test_with_delay(self):
216        retry_ = retry_async.AsyncRetry(
217            predicate=mock.sentinel.predicate,
218            initial=1,
219            maximum=2,
220            multiplier=3,
221            deadline=4,
222            on_error=mock.sentinel.on_error,
223        )
224        new_retry = retry_.with_delay(initial=1, maximum=2, multiplier=3)
225        assert retry_ is not new_retry
226        assert new_retry._initial == 1
227        assert new_retry._maximum == 2
228        assert new_retry._multiplier == 3
229
230        # the rest of the attributes should remain the same
231        assert new_retry._deadline == retry_._deadline
232        assert new_retry._predicate is retry_._predicate
233        assert new_retry._on_error is retry_._on_error
234
235    def test___str__(self):
236        def if_exception_type(exc):
237            return bool(exc)  # pragma: NO COVER
238
239        # Explicitly set all attributes as changed Retry defaults should not
240        # cause this test to start failing.
241        retry_ = retry_async.AsyncRetry(
242            predicate=if_exception_type,
243            initial=1.0,
244            maximum=60.0,
245            multiplier=2.0,
246            deadline=120.0,
247            on_error=None,
248        )
249        assert re.match(
250            (
251                r"<AsyncRetry predicate=<function.*?if_exception_type.*?>, "
252                r"initial=1.0, maximum=60.0, multiplier=2.0, deadline=120.0, "
253                r"on_error=None>"
254            ),
255            str(retry_),
256        )
257
258    @mock.patch("asyncio.sleep", autospec=True)
259    @pytest.mark.asyncio
260    async def test___call___and_execute_success(self, sleep):
261        retry_ = retry_async.AsyncRetry()
262        target = mock.AsyncMock(spec=["__call__"], return_value=42)
263        # __name__ is needed by functools.partial.
264        target.__name__ = "target"
265
266        decorated = retry_(target)
267        target.assert_not_called()
268
269        result = await decorated("meep")
270
271        assert result == 42
272        target.assert_called_once_with("meep")
273        sleep.assert_not_called()
274
275    # Make uniform return half of its maximum, which is the calculated sleep time.
276    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
277    @mock.patch("asyncio.sleep", autospec=True)
278    @pytest.mark.asyncio
279    async def test___call___and_execute_retry(self, sleep, uniform):
280
281        on_error = mock.Mock(spec=["__call__"], side_effect=[None])
282        retry_ = retry_async.AsyncRetry(
283            predicate=retry_async.if_exception_type(ValueError)
284        )
285
286        target = mock.AsyncMock(spec=["__call__"], side_effect=[ValueError(), 42])
287        # __name__ is needed by functools.partial.
288        target.__name__ = "target"
289
290        decorated = retry_(target, on_error=on_error)
291        target.assert_not_called()
292
293        result = await decorated("meep")
294
295        assert result == 42
296        assert target.call_count == 2
297        target.assert_has_calls([mock.call("meep"), mock.call("meep")])
298        sleep.assert_called_once_with(retry_._initial)
299        assert on_error.call_count == 1
300
301    # Make uniform return half of its maximum, which is the calculated sleep time.
302    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
303    @mock.patch("asyncio.sleep", autospec=True)
304    @pytest.mark.asyncio
305    async def test___call___and_execute_retry_hitting_deadline(self, sleep, uniform):
306
307        on_error = mock.Mock(spec=["__call__"], side_effect=[None] * 10)
308        retry_ = retry_async.AsyncRetry(
309            predicate=retry_async.if_exception_type(ValueError),
310            initial=1.0,
311            maximum=1024.0,
312            multiplier=2.0,
313            deadline=9.9,
314        )
315
316        utcnow = datetime.datetime.utcnow()
317        utcnow_patcher = mock.patch(
318            "google.api_core.datetime_helpers.utcnow", return_value=utcnow
319        )
320
321        target = mock.AsyncMock(spec=["__call__"], side_effect=[ValueError()] * 10)
322        # __name__ is needed by functools.partial.
323        target.__name__ = "target"
324
325        decorated = retry_(target, on_error=on_error)
326        target.assert_not_called()
327
328        with utcnow_patcher as patched_utcnow:
329            # Make sure that calls to fake asyncio.sleep() also advance the mocked
330            # time clock.
331            def increase_time(sleep_delay):
332                patched_utcnow.return_value += datetime.timedelta(seconds=sleep_delay)
333
334            sleep.side_effect = increase_time
335
336            with pytest.raises(exceptions.RetryError):
337                await decorated("meep")
338
339        assert target.call_count == 5
340        target.assert_has_calls([mock.call("meep")] * 5)
341        assert on_error.call_count == 5
342
343        # check the delays
344        assert sleep.call_count == 4  # once between each successive target calls
345        last_wait = sleep.call_args.args[0]
346        total_wait = sum(call_args.args[0] for call_args in sleep.call_args_list)
347
348        assert last_wait == 2.9  # and not 8.0, because the last delay was shortened
349        assert total_wait == 9.9  # the same as the deadline
350
351    @mock.patch("asyncio.sleep", autospec=True)
352    @pytest.mark.asyncio
353    async def test___init___without_retry_executed(self, sleep):
354        _some_function = mock.Mock()
355
356        retry_ = retry_async.AsyncRetry(
357            predicate=retry_async.if_exception_type(ValueError), on_error=_some_function
358        )
359        # check the proper creation of the class
360        assert retry_._on_error is _some_function
361
362        target = mock.AsyncMock(spec=["__call__"], side_effect=[42])
363        # __name__ is needed by functools.partial.
364        target.__name__ = "target"
365
366        wrapped = retry_(target)
367
368        result = await wrapped("meep")
369
370        assert result == 42
371        target.assert_called_once_with("meep")
372        sleep.assert_not_called()
373        _some_function.assert_not_called()
374
375    # Make uniform return half of its maximum, which is the calculated sleep time.
376    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
377    @mock.patch("asyncio.sleep", autospec=True)
378    @pytest.mark.asyncio
379    async def test___init___when_retry_is_executed(self, sleep, uniform):
380        _some_function = mock.Mock()
381
382        retry_ = retry_async.AsyncRetry(
383            predicate=retry_async.if_exception_type(ValueError), on_error=_some_function
384        )
385        # check the proper creation of the class
386        assert retry_._on_error is _some_function
387
388        target = mock.AsyncMock(
389            spec=["__call__"], side_effect=[ValueError(), ValueError(), 42]
390        )
391        # __name__ is needed by functools.partial.
392        target.__name__ = "target"
393
394        wrapped = retry_(target)
395        target.assert_not_called()
396
397        result = await wrapped("meep")
398
399        assert result == 42
400        assert target.call_count == 3
401        assert _some_function.call_count == 2
402        target.assert_has_calls([mock.call("meep"), mock.call("meep")])
403        sleep.assert_any_call(retry_._initial)
404