Skip to content

Commit

Permalink
refactor VPC service find_usage()
Browse files Browse the repository at this point in the history
  • Loading branch information
@jantman
jantman committed Jul 1, 2015
1 parent bb5ca36 commit 0d16b37
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 104 deletions.
19 changes: 17 additions & 2 deletions awslimitchecker/services/vpc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,14 +72,25 @@ def find_usage(self):
self.connect()
for lim in self.limits.values():
lim._reset_usage()
self._find_usage_vpcs()
self._find_usage_subnets()
self._find_usage_ACLs()
self._find_usage_route_tables()
self._find_usage_gateways()
self._have_usage = True
logger.debug("Done checking usage.")

def _find_usage_vpcs(self):
"""find usage for VPCs"""
# overall number of VPCs
vpcs = self.conn.get_all_vpcs()
self.limits['VPCs']._add_current_usage(
len(vpcs),
aws_type='AWS::EC2::VPC'
)

def _find_usage_subnets(self):
"""find usage for Subnets"""
# subnets per VPC
subnets = defaultdict(int)
for subnet in self.conn.get_all_subnets():
Expand All @@ -92,6 +103,8 @@ def find_usage(self):
id=vpc_id
)

def _find_usage_ACLs(self):
"""find usage for ACLs"""
# Network ACLs per VPC
acls = defaultdict(int)
for acl in self.conn.get_all_network_acls():
Expand All @@ -110,6 +123,8 @@ def find_usage(self):
id=vpc_id,
)

def _find_usage_route_tables(self):
"""find usage for route tables"""
# Route tables per VPC
tables = defaultdict(int)
for table in self.conn.get_all_route_tables():
Expand All @@ -128,14 +143,14 @@ def find_usage(self):
id=vpc_id,
)

def _find_usage_gateways(self):
"""find usage for Internet Gateways"""
# Internet gateways
gws = self.conn.get_all_internet_gateways()
self.limits['Internet gateways']._add_current_usage(
len(gws),
aws_type='AWS::EC2::InternetGateway',
)
self._have_usage = True
logger.debug("Done checking usage.")

def get_limits(self):
"""
Expand Down
148 changes: 46 additions & 102 deletions awslimitchecker/tests/services/test_vpc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,9 @@
sys.version_info[0] < 3 or
sys.version_info[0] == 3 and sys.version_info[1] < 4
):
from mock import patch, call, Mock
from mock import patch, call, Mock, DEFAULT
else:
from unittest.mock import patch, call, Mock
from unittest.mock import patch, call, Mock, DEFAULT


class Test_VpcService(object):
Expand Down Expand Up @@ -118,40 +118,52 @@ def test_get_limits_again(self):
res = cls.get_limits()
assert res == mock_limits

def test_find_usage_vpc(self):
def test_find_usage(self):
mock_conn = Mock(spec_set=VPCConnection)

with patch('%s.connect' % self.pb) as mock_connect:
with patch.multiple(
self.pb,
_find_usage_vpcs=DEFAULT,
_find_usage_subnets=DEFAULT,
_find_usage_ACLs=DEFAULT,
_find_usage_route_tables=DEFAULT,
_find_usage_gateways=DEFAULT,
) as mocks:
cls = _VpcService(21, 43)
cls.conn = mock_conn
assert cls._have_usage is False
cls.find_usage()
assert mock_connect.mock_calls == [call()]
assert cls._have_usage is True
assert mock_conn.mock_calls == []
for x in [
'_find_usage_vpcs',
'_find_usage_subnets',
'_find_usage_ACLs',
'_find_usage_route_tables',
'_find_usage_gateways',
]:
assert mocks[x].mock_calls == [call()]

def test_find_usage_vpcs(self):
mock1 = Mock(spec_set=VPC)
type(mock1).id = 'vpc-1'
mock2 = Mock(spec_set=VPC)
type(mock2).id = 'vpc-2'

vpcs = [mock1, mock2]
subnets = []
acls = []
tables = []
gateways = []

mock_conn = Mock(spec_set=VPCConnection)
mock_conn.get_all_vpcs.return_value = vpcs
mock_conn.get_all_subnets.return_value = subnets
mock_conn.get_all_network_acls.return_value = acls
mock_conn.get_all_route_tables.return_value = tables
mock_conn.get_all_internet_gateways.return_value = gateways

with patch('%s.connect' % self.pb) as mock_connect:
cls = _VpcService(21, 43)
cls.conn = mock_conn
assert cls._have_usage is False
cls.find_usage()
assert mock_connect.mock_calls == [call()]
assert cls._have_usage is True
cls = _VpcService(21, 43)
cls.conn = mock_conn
cls._find_usage_vpcs()
assert len(cls.limits['VPCs'].get_current_usage()) == 1
assert cls.limits['VPCs'].get_current_usage()[0].get_value() == 2
assert mock_conn.mock_calls == [
call.get_all_vpcs(),
call.get_all_subnets(),
call.get_all_network_acls(),
call.get_all_route_tables(),
call.get_all_internet_gateways(),
]

def test_find_usage_subnets(self):
Expand All @@ -162,39 +174,20 @@ def test_find_usage_subnets(self):
mock3 = Mock(spec_set=Subnet)
type(mock3).vpc_id = 'vpc-2'

vpcs = []
subnets = [mock1, mock2, mock3]
acls = []
tables = []
gateways = []

mock_conn = Mock(spec_set=VPCConnection)
mock_conn.get_all_vpcs.return_value = vpcs
mock_conn.get_all_subnets.return_value = subnets
mock_conn.get_all_network_acls.return_value = acls
mock_conn.get_all_route_tables.return_value = tables
mock_conn.get_all_internet_gateways.return_value = gateways

with patch('%s.connect' % self.pb) as mock_connect:
cls = _VpcService(21, 43)
cls.conn = mock_conn
assert cls._have_usage is False
cls.find_usage()
assert mock_connect.mock_calls == [call()]
assert cls._have_usage is True
assert cls.limits['VPCs'].get_current_usage()[0].get_value() == 0
cls = _VpcService(21, 43)
cls.conn = mock_conn
cls._find_usage_subnets()
usage = sorted(cls.limits['Subnets per VPC'].get_current_usage())
assert len(usage) == 2
assert usage[0].get_value() == 1
assert usage[0].id == 'vpc-2'
assert usage[1].get_value() == 2
assert usage[1].id == 'vpc-1'
assert mock_conn.mock_calls == [
call.get_all_vpcs(),
call.get_all_subnets(),
call.get_all_network_acls(),
call.get_all_route_tables(),
call.get_all_internet_gateways(),
]

def test_find_usage_acls(self):
Expand All @@ -211,26 +204,13 @@ def test_find_usage_acls(self):
type(mock3).vpc_id = 'vpc-2'
type(mock3).network_acl_entries = [1, 2, 3, 4, 5]

vpcs = []
subnets = []
acls = [mock1, mock2, mock3]
tables = []
gateways = []

mock_conn = Mock(spec_set=VPCConnection)
mock_conn.get_all_vpcs.return_value = vpcs
mock_conn.get_all_subnets.return_value = subnets
mock_conn.get_all_network_acls.return_value = acls
mock_conn.get_all_route_tables.return_value = tables
mock_conn.get_all_internet_gateways.return_value = gateways

with patch('%s.connect' % self.pb) as mock_connect:
cls = _VpcService(21, 43)
cls.conn = mock_conn
assert cls._have_usage is False
cls.find_usage()
assert mock_connect.mock_calls == [call()]
assert cls._have_usage is True
cls = _VpcService(21, 43)
cls.conn = mock_conn
cls._find_usage_ACLs()
usage = sorted(cls.limits['Network ACLs per VPC'].get_current_usage())
assert len(usage) == 2
assert usage[0].get_value() == 1
Expand All @@ -247,11 +227,7 @@ def test_find_usage_acls(self):
assert entries[2].id == 'acl-3'
assert entries[2].get_value() == 5
assert mock_conn.mock_calls == [
call.get_all_vpcs(),
call.get_all_subnets(),
call.get_all_network_acls(),
call.get_all_route_tables(),
call.get_all_internet_gateways(),
]

def test_find_usage_route_tables(self):
Expand All @@ -268,26 +244,14 @@ def test_find_usage_route_tables(self):
type(mock3).vpc_id = 'vpc-2'
type(mock3).routes = [1, 2, 3, 4, 5]

vpcs = []
subnets = []
acls = []
tables = [mock1, mock2, mock3]
gateways = []

mock_conn = Mock(spec_set=VPCConnection)
mock_conn.get_all_vpcs.return_value = vpcs
mock_conn.get_all_subnets.return_value = subnets
mock_conn.get_all_network_acls.return_value = acls
mock_conn.get_all_route_tables.return_value = tables
mock_conn.get_all_internet_gateways.return_value = gateways

with patch('%s.connect' % self.pb) as mock_connect:
cls = _VpcService(21, 43)
cls.conn = mock_conn
assert cls._have_usage is False
cls.find_usage()
assert mock_connect.mock_calls == [call()]
assert cls._have_usage is True
cls = _VpcService(21, 43)
cls.conn = mock_conn
cls._find_usage_route_tables()
usage = sorted(cls.limits['Route tables per VPC'].get_current_usage())
assert len(usage) == 2
assert usage[0].get_value() == 1
Expand All @@ -304,11 +268,7 @@ def test_find_usage_route_tables(self):
assert entries[2].id == 'rt-3'
assert entries[2].get_value() == 5
assert mock_conn.mock_calls == [
call.get_all_vpcs(),
call.get_all_subnets(),
call.get_all_network_acls(),
call.get_all_route_tables(),
call.get_all_internet_gateways(),
]

def test_find_usage_internet_gateways(self):
Expand All @@ -317,34 +277,18 @@ def test_find_usage_internet_gateways(self):
mock2 = Mock(spec_set=InternetGateway)
type(mock2).id = 'gw-2'

vpcs = []
subnets = []
acls = []
tables = []
gateways = [mock1, mock2]

mock_conn = Mock(spec_set=VPCConnection)
mock_conn.get_all_vpcs.return_value = vpcs
mock_conn.get_all_subnets.return_value = subnets
mock_conn.get_all_network_acls.return_value = acls
mock_conn.get_all_route_tables.return_value = tables
mock_conn.get_all_internet_gateways.return_value = gateways

with patch('%s.connect' % self.pb) as mock_connect:
cls = _VpcService(21, 43)
cls.conn = mock_conn
assert cls._have_usage is False
cls.find_usage()
assert mock_connect.mock_calls == [call()]
assert cls._have_usage is True
cls = _VpcService(21, 43)
cls.conn = mock_conn
cls._find_usage_gateways()
assert len(cls.limits['Internet gateways'].get_current_usage()) == 1
assert cls.limits['Internet gateways'].get_current_usage()[
0].get_value() == 2
assert mock_conn.mock_calls == [
call.get_all_vpcs(),
call.get_all_subnets(),
call.get_all_network_acls(),
call.get_all_route_tables(),
call.get_all_internet_gateways(),
]

Expand Down

0 comments on commit 0d16b37

Please sign in to comment.