mirror of
https://github.com/optim-enterprises-bv/patroni.git
synced 2026-01-08 08:31:27 +00:00
77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
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())
|