1# Copyright 2017 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 itertools
17import re
18
19import mock
20import pytest
21import requests.exceptions
22
23from google.api_core import exceptions
24from google.api_core import retry
25from google.auth import exceptions as auth_exceptions
26
27
28def test_if_exception_type():
29    predicate = retry.if_exception_type(ValueError)
30
31    assert predicate(ValueError())
32    assert not predicate(TypeError())
33
34
35def test_if_exception_type_multiple():
36    predicate = retry.if_exception_type(ValueError, TypeError)
37
38    assert predicate(ValueError())
39    assert predicate(TypeError())
40    assert not predicate(RuntimeError())
41
42
43def test_if_transient_error():
44    assert retry.if_transient_error(exceptions.InternalServerError(""))
45    assert retry.if_transient_error(exceptions.TooManyRequests(""))
46    assert retry.if_transient_error(exceptions.ServiceUnavailable(""))
47    assert retry.if_transient_error(requests.exceptions.ConnectionError(""))
48    assert retry.if_transient_error(requests.exceptions.ChunkedEncodingError(""))
49    assert retry.if_transient_error(auth_exceptions.TransportError(""))
50    assert not retry.if_transient_error(exceptions.InvalidArgument(""))
51
52
53# Make uniform return half of its maximum, which will be the calculated
54# sleep time.
55@mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
56def test_exponential_sleep_generator_base_2(uniform):
57    gen = retry.exponential_sleep_generator(1, 60, multiplier=2)
58
59    result = list(itertools.islice(gen, 8))
60    assert result == [1, 2, 4, 8, 16, 32, 60, 60]
61
62
63@mock.patch("time.sleep", autospec=True)
64@mock.patch(
65    "google.api_core.datetime_helpers.utcnow",
66    return_value=datetime.datetime.min,
67    autospec=True,
68)
69def test_retry_target_success(utcnow, sleep):
70    predicate = retry.if_exception_type(ValueError)
71    call_count = [0]
72
73    def target():
74        call_count[0] += 1
75        if call_count[0] < 3:
76            raise ValueError()
77        return 42
78
79    result = retry.retry_target(target, predicate, range(10), None)
80
81    assert result == 42
82    assert call_count[0] == 3
83    sleep.assert_has_calls([mock.call(0), mock.call(1)])
84
85
86@mock.patch("time.sleep", autospec=True)
87@mock.patch(
88    "google.api_core.datetime_helpers.utcnow",
89    return_value=datetime.datetime.min,
90    autospec=True,
91)
92def test_retry_target_w_on_error(utcnow, sleep):
93    predicate = retry.if_exception_type(ValueError)
94    call_count = {"target": 0}
95    to_raise = ValueError()
96
97    def target():
98        call_count["target"] += 1
99        if call_count["target"] < 3:
100            raise to_raise
101        return 42
102
103    on_error = mock.Mock()
104
105    result = retry.retry_target(target, predicate, range(10), None, on_error=on_error)
106
107    assert result == 42
108    assert call_count["target"] == 3
109
110    on_error.assert_has_calls([mock.call(to_raise), mock.call(to_raise)])
111    sleep.assert_has_calls([mock.call(0), mock.call(1)])
112
113
114@mock.patch("time.sleep", autospec=True)
115@mock.patch(
116    "google.api_core.datetime_helpers.utcnow",
117    return_value=datetime.datetime.min,
118    autospec=True,
119)
120def test_retry_target_non_retryable_error(utcnow, sleep):
121    predicate = retry.if_exception_type(ValueError)
122    exception = TypeError()
123    target = mock.Mock(side_effect=exception)
124
125    with pytest.raises(TypeError) as exc_info:
126        retry.retry_target(target, predicate, range(10), None)
127
128    assert exc_info.value == exception
129    sleep.assert_not_called()
130
131
132@mock.patch("time.sleep", autospec=True)
133@mock.patch("google.api_core.datetime_helpers.utcnow", autospec=True)
134def test_retry_target_deadline_exceeded(utcnow, sleep):
135    predicate = retry.if_exception_type(ValueError)
136    exception = ValueError("meep")
137    target = mock.Mock(side_effect=exception)
138    # Setup the timeline so that the first call takes 5 seconds but the second
139    # call takes 6, which puts the retry over the deadline.
140    utcnow.side_effect = [
141        # The first call to utcnow establishes the start of the timeline.
142        datetime.datetime.min,
143        datetime.datetime.min + datetime.timedelta(seconds=5),
144        datetime.datetime.min + datetime.timedelta(seconds=11),
145    ]
146
147    with pytest.raises(exceptions.RetryError) as exc_info:
148        retry.retry_target(target, predicate, range(10), deadline=10)
149
150    assert exc_info.value.cause == exception
151    assert exc_info.match("Deadline of 10.0s exceeded")
152    assert exc_info.match("last exception: meep")
153    assert target.call_count == 2
154
155
156def test_retry_target_bad_sleep_generator():
157    with pytest.raises(ValueError, match="Sleep generator"):
158        retry.retry_target(mock.sentinel.target, mock.sentinel.predicate, [], None)
159
160
161class TestRetry(object):
162    def test_constructor_defaults(self):
163        retry_ = retry.Retry()
164        assert retry_._predicate == retry.if_transient_error
165        assert retry_._initial == 1
166        assert retry_._maximum == 60
167        assert retry_._multiplier == 2
168        assert retry_._deadline == 120
169        assert retry_._on_error is None
170        assert retry_.deadline == 120
171
172    def test_constructor_options(self):
173        _some_function = mock.Mock()
174
175        retry_ = retry.Retry(
176            predicate=mock.sentinel.predicate,
177            initial=1,
178            maximum=2,
179            multiplier=3,
180            deadline=4,
181            on_error=_some_function,
182        )
183        assert retry_._predicate == mock.sentinel.predicate
184        assert retry_._initial == 1
185        assert retry_._maximum == 2
186        assert retry_._multiplier == 3
187        assert retry_._deadline == 4
188        assert retry_._on_error is _some_function
189
190    def test_with_deadline(self):
191        retry_ = retry.Retry(
192            predicate=mock.sentinel.predicate,
193            initial=1,
194            maximum=2,
195            multiplier=3,
196            deadline=4,
197            on_error=mock.sentinel.on_error,
198        )
199        new_retry = retry_.with_deadline(42)
200        assert retry_ is not new_retry
201        assert new_retry._deadline == 42
202
203        # the rest of the attributes should remain the same
204        assert new_retry._predicate is retry_._predicate
205        assert new_retry._initial == retry_._initial
206        assert new_retry._maximum == retry_._maximum
207        assert new_retry._multiplier == retry_._multiplier
208        assert new_retry._on_error is retry_._on_error
209
210    def test_with_predicate(self):
211        retry_ = retry.Retry(
212            predicate=mock.sentinel.predicate,
213            initial=1,
214            maximum=2,
215            multiplier=3,
216            deadline=4,
217            on_error=mock.sentinel.on_error,
218        )
219        new_retry = retry_.with_predicate(mock.sentinel.predicate)
220        assert retry_ is not new_retry
221        assert new_retry._predicate == mock.sentinel.predicate
222
223        # the rest of the attributes should remain the same
224        assert new_retry._deadline == retry_._deadline
225        assert new_retry._initial == retry_._initial
226        assert new_retry._maximum == retry_._maximum
227        assert new_retry._multiplier == retry_._multiplier
228        assert new_retry._on_error is retry_._on_error
229
230    def test_with_delay_noop(self):
231        retry_ = retry.Retry(
232            predicate=mock.sentinel.predicate,
233            initial=1,
234            maximum=2,
235            multiplier=3,
236            deadline=4,
237            on_error=mock.sentinel.on_error,
238        )
239        new_retry = retry_.with_delay()
240        assert retry_ is not new_retry
241        assert new_retry._initial == retry_._initial
242        assert new_retry._maximum == retry_._maximum
243        assert new_retry._multiplier == retry_._multiplier
244
245    def test_with_delay(self):
246        retry_ = retry.Retry(
247            predicate=mock.sentinel.predicate,
248            initial=1,
249            maximum=2,
250            multiplier=3,
251            deadline=4,
252            on_error=mock.sentinel.on_error,
253        )
254        new_retry = retry_.with_delay(initial=5, maximum=6, multiplier=7)
255        assert retry_ is not new_retry
256        assert new_retry._initial == 5
257        assert new_retry._maximum == 6
258        assert new_retry._multiplier == 7
259
260        # the rest of the attributes should remain the same
261        assert new_retry._deadline == retry_._deadline
262        assert new_retry._predicate is retry_._predicate
263        assert new_retry._on_error is retry_._on_error
264
265    def test_with_delay_partial_options(self):
266        retry_ = retry.Retry(
267            predicate=mock.sentinel.predicate,
268            initial=1,
269            maximum=2,
270            multiplier=3,
271            deadline=4,
272            on_error=mock.sentinel.on_error,
273        )
274        new_retry = retry_.with_delay(initial=4)
275        assert retry_ is not new_retry
276        assert new_retry._initial == 4
277        assert new_retry._maximum == 2
278        assert new_retry._multiplier == 3
279
280        new_retry = retry_.with_delay(maximum=4)
281        assert retry_ is not new_retry
282        assert new_retry._initial == 1
283        assert new_retry._maximum == 4
284        assert new_retry._multiplier == 3
285
286        new_retry = retry_.with_delay(multiplier=4)
287        assert retry_ is not new_retry
288        assert new_retry._initial == 1
289        assert new_retry._maximum == 2
290        assert new_retry._multiplier == 4
291
292        # the rest of the attributes should remain the same
293        assert new_retry._deadline == retry_._deadline
294        assert new_retry._predicate is retry_._predicate
295        assert new_retry._on_error is retry_._on_error
296
297    def test___str__(self):
298        def if_exception_type(exc):
299            return bool(exc)  # pragma: NO COVER
300
301        # Explicitly set all attributes as changed Retry defaults should not
302        # cause this test to start failing.
303        retry_ = retry.Retry(
304            predicate=if_exception_type,
305            initial=1.0,
306            maximum=60.0,
307            multiplier=2.0,
308            deadline=120.0,
309            on_error=None,
310        )
311        assert re.match(
312            (
313                r"<Retry predicate=<function.*?if_exception_type.*?>, "
314                r"initial=1.0, maximum=60.0, multiplier=2.0, deadline=120.0, "
315                r"on_error=None>"
316            ),
317            str(retry_),
318        )
319
320    @mock.patch("time.sleep", autospec=True)
321    def test___call___and_execute_success(self, sleep):
322        retry_ = retry.Retry()
323        target = mock.Mock(spec=["__call__"], return_value=42)
324        # __name__ is needed by functools.partial.
325        target.__name__ = "target"
326
327        decorated = retry_(target)
328        target.assert_not_called()
329
330        result = decorated("meep")
331
332        assert result == 42
333        target.assert_called_once_with("meep")
334        sleep.assert_not_called()
335
336    # Make uniform return half of its maximum, which is the calculated sleep time.
337    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
338    @mock.patch("time.sleep", autospec=True)
339    def test___call___and_execute_retry(self, sleep, uniform):
340
341        on_error = mock.Mock(spec=["__call__"], side_effect=[None])
342        retry_ = retry.Retry(predicate=retry.if_exception_type(ValueError))
343
344        target = mock.Mock(spec=["__call__"], side_effect=[ValueError(), 42])
345        # __name__ is needed by functools.partial.
346        target.__name__ = "target"
347
348        decorated = retry_(target, on_error=on_error)
349        target.assert_not_called()
350
351        result = decorated("meep")
352
353        assert result == 42
354        assert target.call_count == 2
355        target.assert_has_calls([mock.call("meep"), mock.call("meep")])
356        sleep.assert_called_once_with(retry_._initial)
357        assert on_error.call_count == 1
358
359    # Make uniform return half of its maximum, which is the calculated sleep time.
360    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
361    @mock.patch("time.sleep", autospec=True)
362    def test___call___and_execute_retry_hitting_deadline(self, sleep, uniform):
363
364        on_error = mock.Mock(spec=["__call__"], side_effect=[None] * 10)
365        retry_ = retry.Retry(
366            predicate=retry.if_exception_type(ValueError),
367            initial=1.0,
368            maximum=1024.0,
369            multiplier=2.0,
370            deadline=9.9,
371        )
372
373        utcnow = datetime.datetime.utcnow()
374        utcnow_patcher = mock.patch(
375            "google.api_core.datetime_helpers.utcnow", return_value=utcnow
376        )
377
378        target = mock.Mock(spec=["__call__"], side_effect=[ValueError()] * 10)
379        # __name__ is needed by functools.partial.
380        target.__name__ = "target"
381
382        decorated = retry_(target, on_error=on_error)
383        target.assert_not_called()
384
385        with utcnow_patcher as patched_utcnow:
386            # Make sure that calls to fake time.sleep() also advance the mocked
387            # time clock.
388            def increase_time(sleep_delay):
389                patched_utcnow.return_value += datetime.timedelta(seconds=sleep_delay)
390
391            sleep.side_effect = increase_time
392
393            with pytest.raises(exceptions.RetryError):
394                decorated("meep")
395
396        assert target.call_count == 5
397        target.assert_has_calls([mock.call("meep")] * 5)
398        assert on_error.call_count == 5
399
400        # check the delays
401        assert sleep.call_count == 4  # once between each successive target calls
402        last_wait = sleep.call_args.args[0]
403        total_wait = sum(call_args.args[0] for call_args in sleep.call_args_list)
404
405        assert last_wait == 2.9  # and not 8.0, because the last delay was shortened
406        assert total_wait == 9.9  # the same as the deadline
407
408    @mock.patch("time.sleep", autospec=True)
409    def test___init___without_retry_executed(self, sleep):
410        _some_function = mock.Mock()
411
412        retry_ = retry.Retry(
413            predicate=retry.if_exception_type(ValueError), on_error=_some_function
414        )
415        # check the proper creation of the class
416        assert retry_._on_error is _some_function
417
418        target = mock.Mock(spec=["__call__"], side_effect=[42])
419        # __name__ is needed by functools.partial.
420        target.__name__ = "target"
421
422        wrapped = retry_(target)
423
424        result = wrapped("meep")
425
426        assert result == 42
427        target.assert_called_once_with("meep")
428        sleep.assert_not_called()
429        _some_function.assert_not_called()
430
431    # Make uniform return half of its maximum, which is the calculated sleep time.
432    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
433    @mock.patch("time.sleep", autospec=True)
434    def test___init___when_retry_is_executed(self, sleep, uniform):
435        _some_function = mock.Mock()
436
437        retry_ = retry.Retry(
438            predicate=retry.if_exception_type(ValueError), on_error=_some_function
439        )
440        # check the proper creation of the class
441        assert retry_._on_error is _some_function
442
443        target = mock.Mock(
444            spec=["__call__"], side_effect=[ValueError(), ValueError(), 42]
445        )
446        # __name__ is needed by functools.partial.
447        target.__name__ = "target"
448
449        wrapped = retry_(target)
450        target.assert_not_called()
451
452        result = wrapped("meep")
453
454        assert result == 42
455        assert target.call_count == 3
456        assert _some_function.call_count == 2
457        target.assert_has_calls([mock.call("meep"), mock.call("meep")])
458        sleep.assert_any_call(retry_._initial)
459