don't use table of test cases and reduce duplication in get_policy method

This commit is contained in:
Brady 2020-01-27 10:35:46 -05:00
parent 2a2ff32dec
commit cb07d17e13
2 changed files with 35 additions and 100 deletions

View File

@ -14,22 +14,10 @@ class Policy:
self.statements = [] self.statements = []
self.parent = parent self.parent = parent
def __repr__(self):
return json.dumps(self.get_policy())
def wire_format(self): def wire_format(self):
return json.dumps( p = self.get_policy()
{ p["Policy"] = json.dumps(p["Policy"])
"Policy": json.dumps( return json.dumps(p)
{
"Version": "2012-10-17",
"Id": "default",
"Statement": self.statements,
}
),
"RevisionId": self.revision,
}
)
def get_policy(self): def get_policy(self):
return { return {
@ -81,6 +69,7 @@ class Policy:
# transform field names and values # transform field names and values
self.transform_property(obj, "StatementId", "Sid", self.nop_formatter) self.transform_property(obj, "StatementId", "Sid", self.nop_formatter)
self.transform_property(obj, "Principal", "Principal", self.principal_formatter) self.transform_property(obj, "Principal", "Principal", self.principal_formatter)
self.transform_property( self.transform_property(
obj, "SourceArn", "SourceArn", self.source_arn_formatter obj, "SourceArn", "SourceArn", self.source_arn_formatter
) )

View File

@ -12,93 +12,39 @@ class MockLambdaFunction:
self.policy = None self.policy = None
class TC:
def __init__(self, lambda_arn, statement, expected):
self.statement = statement
self.expected = expected
self.fn = MockLambdaFunction(lambda_arn)
self.policy = Policy(self.fn)
def Run(self, parent):
self.policy.add_statement(json.dumps(self.statement))
parent.assertDictEqual(self.expected, self.policy.statements[0])
sid = self.statement.get("StatementId", None)
if sid == None:
raise "TestCase.statement does not contain StatementId"
self.policy.del_statement(sid)
parent.assertEqual([], self.policy.statements)
class TestPolicy(unittest.TestCase): class TestPolicy(unittest.TestCase):
def test(self): def test_policy(self):
tt = [ policy = Policy(MockLambdaFunction("arn"))
TC( statement = {
# lambda_arn
"arn",
{ # statement
"StatementId": "statement0", "StatementId": "statement0",
"Action": "lambda:InvokeFunction", "Action": "lambda:InvokeFunction",
"FunctionName": "function_name", "FunctionName": "function_name",
"Principal": "events.amazonaws.com", "Principal": "events.amazonaws.com",
}, "SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name",
{ # expected "SourceAccount": "111111111111",
}
expected = {
"Action": "lambda:InvokeFunction", "Action": "lambda:InvokeFunction",
"FunctionName": "function_name", "FunctionName": "function_name",
"Principal": {"Service": "events.amazonaws.com"}, "Principal": {"Service": "events.amazonaws.com"},
"Effect": "Allow", "Effect": "Allow",
"Resource": "arn:$LATEST", "Resource": "arn:$LATEST",
"Sid": "statement0", "Sid": "statement0",
},
),
TC(
# lambda_arn
"arn",
{ # statement
"StatementId": "statement1",
"Action": "lambda:InvokeFunction",
"FunctionName": "function_name",
"Principal": "events.amazonaws.com",
"SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name",
},
{
"Action": "lambda:InvokeFunction",
"FunctionName": "function_name",
"Principal": {"Service": "events.amazonaws.com"},
"Effect": "Allow",
"Resource": "arn:$LATEST",
"Sid": "statement1",
"Condition": { "Condition": {
"ArnLike": { "ArnLike": {
"AWS:SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name" "AWS:SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name",
},
"StringEquals": {"AWS:SourceAccount": "111111111111"},
},
} }
},
},
),
TC(
# lambda_arn
"arn",
{ # statement
"StatementId": "statement2",
"Action": "lambda:InvokeFunction",
"FunctionName": "function_name",
"Principal": "events.amazonaws.com",
"SourceAccount": "111111111111",
},
{ # expected
"Action": "lambda:InvokeFunction",
"FunctionName": "function_name",
"Principal": {"Service": "events.amazonaws.com"},
"Effect": "Allow",
"Resource": "arn:$LATEST",
"Sid": "statement2",
"Condition": {
"StringEquals": {"AWS:SourceAccount": "111111111111"}
},
},
),
]
for tc in tt: policy.add_statement(json.dumps(statement))
tc.Run(self) self.assertDictEqual(expected, policy.statements[0])
sid = statement.get("StatementId", None)
if sid == None:
raise "TestCase.statement does not contain StatementId"
policy.del_statement(sid)
self.assertEqual([], policy.statements)