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
|
import sys
import unittest
from collections import namedtuple
from unittest.mock import Mock, patch, PropertyMock
import botocore
import botocore.awsrequest
from patroni.scripts.aws import AWSConnection, main as _main
class MockVolumes(object):
@staticmethod
def filter(*args, **kwargs):
oid = namedtuple('Volume', 'id')
return [oid(id='a'), oid(id='b')]
class MockEc2Connection(object):
volumes = MockVolumes()
@staticmethod
def create_tags(Resources, **kwargs):
if len(Resources) == 0:
raise botocore.exceptions.ClientError({'Error': {'Code': 503, 'Message': 'Request limit exceeded'}},
'create_tags')
return True
class MockIMDSFetcher(object):
def __init__(self, timeout):
pass
@staticmethod
def _fetch_metadata_token():
return ''
@staticmethod
def _get_request(*args):
return botocore.awsrequest.AWSResponse(url='', status_code=200, headers={}, raw=None)
@patch('boto3.resource', Mock(return_value=MockEc2Connection()))
@patch('patroni.scripts.aws.IMDSFetcher', MockIMDSFetcher)
class TestAWSConnection(unittest.TestCase):
@patch.object(botocore.awsrequest.AWSResponse, 'text',
PropertyMock(return_value='{"instanceId": "012345", "region": "eu-west-1"}'))
def test_on_role_change(self):
conn = AWSConnection('test')
self.assertTrue(conn.on_role_change('primary'))
with patch.object(MockVolumes, 'filter', Mock(return_value=[])):
conn._retry.max_tries = 1
self.assertFalse(conn.on_role_change('primary'))
@patch.object(MockIMDSFetcher, '_get_request', Mock(side_effect=Exception('foo')))
def test_non_aws(self):
conn = AWSConnection('test')
self.assertFalse(conn.on_role_change("primary"))
@patch.object(botocore.awsrequest.AWSResponse, 'text', PropertyMock(return_value='boo'))
def test_aws_bizarre_response(self):
conn = AWSConnection('test')
self.assertFalse(conn.aws_available())
@patch.object(MockIMDSFetcher, '_get_request', Mock(return_value=botocore.awsrequest.AWSResponse(
url='', status_code=503, headers={}, raw=None)))
@patch('sys.exit', Mock())
def test_main(self):
self.assertIsNone(_main())
sys.argv = ['aws.py', 'on_start', 'replica', 'foo']
self.assertIsNone(_main())
|
