1from tests import mock
2from tests import unittest
3
4from botocore.awsrequest import AWSResponse
5from botocore.retries import standard, special
6
7
8def create_fake_op_model(service_name):
9    model = mock.Mock()
10    model.service_model.service_name = service_name
11    return model
12
13
14class TestRetryIDPCommunicationError(unittest.TestCase):
15    def setUp(self):
16        self.checker = special.RetryIDPCommunicationError()
17
18    def test_only_retries_error_for_sts(self):
19        context = standard.RetryContext(
20            attempt_number=1, operation_model=create_fake_op_model('s3'),
21            parsed_response={
22                'Error': {'Code': 'IDPCommunicationError',
23                          'Message': 'message'}},
24            http_response=AWSResponse(
25                status_code=400, raw=None, headers={},
26                url='https://foo'),
27            caught_exception=None, )
28        self.assertFalse(self.checker.is_retryable(context))
29
30    def test_can_retry_idp_communication_error(self):
31        context = standard.RetryContext(
32            attempt_number=1, operation_model=create_fake_op_model('sts'),
33            parsed_response={
34                'Error': {'Code': 'IDPCommunicationError',
35                          'Message': 'message'}},
36            http_response=AWSResponse(
37                status_code=400, raw=None, headers={},
38                url='https://foo'),
39            caught_exception=None, )
40        self.assertTrue(self.checker.is_retryable(context))
41
42    def test_not_idp_communication_error(self):
43        context = standard.RetryContext(
44            attempt_number=1, operation_model=create_fake_op_model('sts'),
45            parsed_response={
46                'Error': {'Code': 'NotIDPCommunicationError',
47                          'Message': 'message'}},
48            http_response=AWSResponse(
49                status_code=400, raw=None, headers={},
50                url='https://foo'),
51            caught_exception=None, )
52        self.assertFalse(self.checker.is_retryable(context))
53
54
55class TestRetryDDBChecksumError(unittest.TestCase):
56    def setUp(self):
57        self.checker = special.RetryDDBChecksumError()
58
59    def raw_stream(self, contents):
60        raw = mock.Mock()
61        raw.stream.return_value = [contents]
62        return raw
63
64    def test_checksum_not_in_header(self):
65        context = standard.RetryContext(
66            attempt_number=1, operation_model=create_fake_op_model('dynamodb'),
67            parsed_response={
68                'Anything': ["foo"],
69            },
70            http_response=AWSResponse(
71                status_code=200, raw=self.raw_stream(b'foo'),
72                headers={},
73                url='https://foo'),
74            caught_exception=None,
75        )
76        self.assertFalse(self.checker.is_retryable(context))
77
78    def test_checksum_matches(self):
79        context = standard.RetryContext(
80            attempt_number=1, operation_model=create_fake_op_model('dynamodb'),
81            parsed_response={
82                'Anything': ["foo"],
83            },
84            http_response=AWSResponse(
85                status_code=200, raw=self.raw_stream(b'foo'),
86                headers={'x-amz-crc32': '2356372769'},
87                url='https://foo'),
88            caught_exception=None
89        )
90        self.assertFalse(self.checker.is_retryable(context))
91
92    def test_checksum_not_matches(self):
93        context = standard.RetryContext(
94            attempt_number=1, operation_model=create_fake_op_model('dynamodb'),
95            parsed_response={
96                'Anything': ["foo"],
97            },
98            http_response=AWSResponse(
99                status_code=200, raw=self.raw_stream(b'foo'),
100                headers={'x-amz-crc32': '2356372768'},
101                url='https://foo'),
102            caught_exception=None
103        )
104        self.assertTrue(self.checker.is_retryable(context))
105
106    def test_checksum_check_only_for_dynamodb(self):
107        context = standard.RetryContext(
108            attempt_number=1, operation_model=create_fake_op_model('s3'),
109            parsed_response={
110                'Anything': ["foo"],
111            },
112            http_response=AWSResponse(
113                status_code=200, raw=self.raw_stream(b'foo'),
114                headers={'x-amz-crc32': '2356372768'},
115                url='https://foo'),
116            caught_exception=None
117        )
118        self.assertFalse(self.checker.is_retryable(context))
119