From cb07d17e1388cce91abe7e8c9de55d04af58f1f8 Mon Sep 17 00:00:00 2001
From: Brady <brady.jubic@gsa.gov>
Date: Mon, 27 Jan 2020 10:35:46 -0500
Subject: [PATCH] don't use table of test cases and reduce duplication in
 get_policy method

---
 moto/awslambda/policy.py            |  19 +----
 tests/test_awslambda/test_policy.py | 116 ++++++++--------------------
 2 files changed, 35 insertions(+), 100 deletions(-)

diff --git a/moto/awslambda/policy.py b/moto/awslambda/policy.py
index 66ec728f..495e2cee 100644
--- a/moto/awslambda/policy.py
+++ b/moto/awslambda/policy.py
@@ -14,22 +14,10 @@ class Policy:
         self.statements = []
         self.parent = parent
 
-    def __repr__(self):
-        return json.dumps(self.get_policy())
-
     def wire_format(self):
-        return json.dumps(
-            {
-                "Policy": json.dumps(
-                    {
-                        "Version": "2012-10-17",
-                        "Id": "default",
-                        "Statement": self.statements,
-                    }
-                ),
-                "RevisionId": self.revision,
-            }
-        )
+        p = self.get_policy()
+        p["Policy"] = json.dumps(p["Policy"])
+        return json.dumps(p)
 
     def get_policy(self):
         return {
@@ -81,6 +69,7 @@ class Policy:
         # transform field names and values
         self.transform_property(obj, "StatementId", "Sid", self.nop_formatter)
         self.transform_property(obj, "Principal", "Principal", self.principal_formatter)
+
         self.transform_property(
             obj, "SourceArn", "SourceArn", self.source_arn_formatter
         )
diff --git a/tests/test_awslambda/test_policy.py b/tests/test_awslambda/test_policy.py
index 2571926e..d3fe3b0b 100644
--- a/tests/test_awslambda/test_policy.py
+++ b/tests/test_awslambda/test_policy.py
@@ -12,93 +12,39 @@ class MockLambdaFunction:
         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)
+class TestPolicy(unittest.TestCase):
+    def test_policy(self):
+        policy = Policy(MockLambdaFunction("arn"))
+        statement = {
+            "StatementId": "statement0",
+            "Action": "lambda:InvokeFunction",
+            "FunctionName": "function_name",
+            "Principal": "events.amazonaws.com",
+            "SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name",
+            "SourceAccount": "111111111111",
+        }
 
-    def Run(self, parent):
-        self.policy.add_statement(json.dumps(self.statement))
-        parent.assertDictEqual(self.expected, self.policy.statements[0])
+        expected = {
+            "Action": "lambda:InvokeFunction",
+            "FunctionName": "function_name",
+            "Principal": {"Service": "events.amazonaws.com"},
+            "Effect": "Allow",
+            "Resource": "arn:$LATEST",
+            "Sid": "statement0",
+            "Condition": {
+                "ArnLike": {
+                    "AWS:SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name",
+                },
+                "StringEquals": {"AWS:SourceAccount": "111111111111"},
+            },
+        }
 
-        sid = self.statement.get("StatementId", None)
+        policy.add_statement(json.dumps(statement))
+        self.assertDictEqual(expected, policy.statements[0])
+
+        sid = 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):
-    def test(self):
-        tt = [
-            TC(
-                # lambda_arn
-                "arn",
-                {  # statement
-                    "StatementId": "statement0",
-                    "Action": "lambda:InvokeFunction",
-                    "FunctionName": "function_name",
-                    "Principal": "events.amazonaws.com",
-                },
-                {  # expected
-                    "Action": "lambda:InvokeFunction",
-                    "FunctionName": "function_name",
-                    "Principal": {"Service": "events.amazonaws.com"},
-                    "Effect": "Allow",
-                    "Resource": "arn:$LATEST",
-                    "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": {
-                        "ArnLike": {
-                            "AWS:SourceArn": "arn:aws:events:us-east-1:111111111111:rule/rule_name"
-                        }
-                    },
-                },
-            ),
-            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:
-            tc.Run(self)
+        policy.del_statement(sid)
+        self.assertEqual([], policy.statements)