diff --git a/ec2_rule_transformer.py b/ec2_rule_transformer.py index a83c63c..5f32a3a 100644 --- a/ec2_rule_transformer.py +++ b/ec2_rule_transformer.py @@ -8,8 +8,10 @@ class EC2RuleTransformer: self.ec2_connection = ec2_connection def to_rule(self, ec2_rule): - rule_args = deepcopy(vars(ec2_rule)) - self._delete_unused_rule_args(rule_args) + rule_args = {} + rule_args['ip_protocol'] = ec2_rule.ip_protocol + rule_args['from_port'] = ec2_rule.from_port + rule_args['to_port'] = ec2_rule.to_port if ec2_rule.grants[0].cidr_ip: rule_args['ip_range'] = ec2_rule.grants[0].cidr_ip @@ -17,11 +19,4 @@ class EC2RuleTransformer: group_id = ec2_rule.grants[0].group_id rule_args['group_name'] = self.ec2_connection.get_all_security_groups(group_ids=group_id)[0].name - return Rule(**rule_args) - - def _delete_unused_rule_args(self, rule_args): - del rule_args['grants'] - del rule_args['parent'] - del rule_args['item'] - del rule_args['ipRanges'] - del rule_args['groups'] \ No newline at end of file + return Rule(**rule_args) \ No newline at end of file diff --git a/openstack_rule_transformer.py b/openstack_rule_transformer.py index 5a84cb9..300198a 100644 --- a/openstack_rule_transformer.py +++ b/openstack_rule_transformer.py @@ -3,11 +3,11 @@ from rule import Rule class OpenstackRuleTransformer: - def to_rule(self, openstack_rule): - rule_args = deepcopy(openstack_rule) - - self._delete_unused_rule_args(rule_args) + rule_args = {} + rule_args['ip_protocol'] = openstack_rule['ip_protocol'] + rule_args['from_port'] = str(openstack_rule['from_port']) + rule_args['to_port'] = str(openstack_rule['to_port']) if 'cidr' in openstack_rule['ip_range']: rule_args['ip_range'] = openstack_rule['ip_range']['cidr'] @@ -15,8 +15,3 @@ class OpenstackRuleTransformer: rule_args['group_name'] = openstack_rule['group']['name'] return Rule(**rule_args) - - def _delete_unused_rule_args(self, rule_args): - del rule_args['group'] - del rule_args['parent_group_id'] - del rule_args['id'] \ No newline at end of file diff --git a/tests/unit/test_openstack_rule_transformer.py b/tests/unit/test_openstack_rule_transformer.py index 887a3b2..b86466c 100644 --- a/tests/unit/test_openstack_rule_transformer.py +++ b/tests/unit/test_openstack_rule_transformer.py @@ -8,7 +8,7 @@ class TestOpenstackRuleTransformer(unittest.TestCase): def setUp(self): self.openstack_rule_transformer = OpenstackRuleTransformer() - def test_should_copy_to_port(self): + def test_should_copy_to_port_as_string(self): openstack_rule = { 'ip_protocol': 'abc', 'from_port': 123, @@ -19,4 +19,17 @@ class TestOpenstackRuleTransformer(unittest.TestCase): 'id': 18 } rule = self.openstack_rule_transformer.to_rule(openstack_rule) - self.assertEqual(rule.to_port, openstack_rule['to_port']) \ No newline at end of file + self.assertEqual(rule.to_port, str(openstack_rule['to_port'])) + + def test_should_copy_from_port_as_string(self): + openstack_rule = { + 'ip_protocol': 'abc', + 'from_port': 123, + 'to_port': 456, + 'group': {}, + 'parent_group_id': 55, + 'ip_range': {'cidr': '9.8.7.6/55'}, + 'id': 18 + } + rule = self.openstack_rule_transformer.to_rule(openstack_rule) + self.assertEqual(rule.from_port, str(openstack_rule['from_port'])) \ No newline at end of file