diff --git a/helpers/aws.py b/helpers/aws.py index d09da2e0..3b72c492 100644 --- a/helpers/aws.py +++ b/helpers/aws.py @@ -46,10 +46,8 @@ class AWSConnection: tags = {'Name': 'spilo_'+self.cluster_name, 'Role': role, 'Instance': self.instance_id} try: conn = boto.ec2.connect_to_region(self.region) - # get all volumes attached to the current instance volumes = conn.get_all_volumes(filters={'attachment.instance-id': self.instance_id}) - if volumes: - conn.create_tags([v.id for v in volumes], tags) + conn.create_tags([v.id for v in volumes], tags) except Exception as e: logger.info('could not set tags for EBS storage devices attached: {}'.format(e)) return False @@ -69,45 +67,5 @@ class AWSConnection: return True def on_role_change(self, new_role): - self._tag_ec2(new_role) - self._tag_ebs(new_role) - -if __name__ == '__main__': - import yaml - config_string = """ -loop_wait: 10 -restapi: - listen: 0.0.0.0:8008 - connect_address: 127.0.0.1:5432 -etcd: - scope: test - ttl: 30 - host: 127.0.0.1:8080 -postgresql: - name: postgresql_foo - listen: 0.0.0.0:5432 - connect_address: 127.0.0.1:5432 - data_dir: /home/postgres/pgdata/data - replication: - username: standby - password: standby - network: 0.0.0.0/0 - superuser: - password: zalando - admin: - username: admin - password: admin - parameters: - archive_mode: "on" - wal_level: hot_standby - max_wal_senders: 5 - wal_keep_segments: 8 - archive_timeout: 1800s - max_replication_slots: 5 - hot_standby: "on" - ssl: "on" -""" - awsconnection = AWSConnection(yaml.load(config_string)) - print "AWS available: {}, Cluster_name: {}".format(awsconnection.available, awsconnection.cluster_name) - if awsconnection.available: - print "AWS Region: {}, Instance_id: {}".format(awsconnection.region, awsconnection.instance_id) + ret = self._tag_ec2(new_role) + return self._tag_ebs(new_role) and ret diff --git a/tests/test_aws.py b/tests/test_aws.py new file mode 100644 index 00000000..e350a5c1 --- /dev/null +++ b/tests/test_aws.py @@ -0,0 +1,112 @@ +import unittest +import requests +import boto.ec2 +from collections import namedtuple +from helpers.aws import AWSConnection +from requests.exceptions import RequestException +import yaml + + +class MockEc2Connection: + + def __init__(self, error=False): + self.error = error + + def get_all_volumes(self, filters): + if self.error: + raise Exception("get_all_volumes") + oid = namedtuple('Volume', 'id') + return [oid(id='a'), oid(id='b')] + + def create_tags(self, objects, tags): + if self.error or len(objects) == 0: + raise Exception("create_tags") + return True + + +class TestAWSConnection(unittest.TestCase): + + def __init__(self, method_name='runTest'): + super(TestAWSConnection, self).__init__(method_name) + + def set_error(self): + self.error = True + + def set_ok(self): + self.error = False + + def boto_ec2_connect_to_region(self, region): + return MockEc2Connection(self.error) + + def requests_get(self, url, **kwargs): + if self.error: + raise RequestException("foo") + result = namedtuple('Request', 'ok content') + result.ok = True + if url.split('/')[-1] == 'availability-zone': + result.content = 'eu-west-1c' + else: + result.content = 'foo' + return result + + def setUp(self): + self.error = False + requests.get = self.requests_get + boto.ec2.connect_to_region = self.boto_ec2_connect_to_region + self.config_string = """ +loop_wait: 10 +restapi: + listen: 0.0.0.0:8008 + connect_address: 127.0.0.1:5432 +etcd: + scope: test + ttl: 30 + host: 127.0.0.1:8080 +postgresql: + name: postgresql_foo + listen: 0.0.0.0:5432 + connect_address: 127.0.0.1:5432 + data_dir: /home/postgres/pgdata/data + replication: + username: standby + password: standby + network: 0.0.0.0/0 + superuser: + password: zalando + admin: + username: admin + password: admin + parameters: + archive_mode: "on" + wal_level: hot_standby + max_wal_senders: 5 + wal_keep_segments: 8 + archive_timeout: 1800s + max_replication_slots: 5 + hot_standby: "on" + ssl: "on" +""" + self.conn = AWSConnection(yaml.load(self.config_string)) + + def test_aws_available(self): + self.assertTrue(self.conn.aws_available()) + + def test_on_role_change(self): + self.assertTrue(self.conn._tag_ebs('master')) + self.assertTrue(self.conn._tag_ec2('master')) + self.assertTrue(self.conn.on_role_change('master')) + + def test_non_aws(self): + self.set_error() + conn = AWSConnection(yaml.load(self.config_string)) + self.assertFalse(conn.aws_available()) + self.assertFalse(conn._tag_ebs('master')) + self.assertFalse(conn._tag_ec2('master')) + + def test_aws_tag_ebs_error(self): + self.set_error() + self.assertFalse(self.conn._tag_ebs("master")) + + def test_aws_tag_ec2_error(self): + self.set_error() + self.assertFalse(self.conn._tag_ec2("master"))