From 1f633f421291ac04adac1b4844c5f03885fca0a5 Mon Sep 17 00:00:00 2001
From: Daniel Milnes <thebeanogamer@gmail.com>
Date: Wed, 8 Jan 2020 21:33:11 +0000
Subject: [PATCH] Fix timed emails (#16)

* Fix keyerror

* Add stage environment variable

* Shrink timeout

* Remove event from connect()

* Fix for loop

* Properly fix for loop

* Remove event
---
 linkshortener/notifications.py | 10 +++++-----
 linkshortener/shortener.py     | 14 +++++++-------
 serverless.yml                 |  3 ++-
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/linkshortener/notifications.py b/linkshortener/notifications.py
index 74cfd97..cf5b0e7 100644
--- a/linkshortener/notifications.py
+++ b/linkshortener/notifications.py
@@ -7,8 +7,8 @@
 from linkshortener.shortener import connect
 
 
-def generate(event):
-    db = connect(event)
+def generate():
+    db = connect()
     page = (
         jinja2.Environment(
             loader=jinja2.FileSystemLoader(searchpath="./linkshortener/templates/"),
@@ -37,10 +37,10 @@ def view(event, context):
 
 
 def summary(event, context):
-    if event["httpMethod"] is None:
+    if event.get("httpMethod") is None:
         db = connect()
         new = False
-        for i in db.scan():
+        for i in db.scan()["Items"]:
             if i["uses"]["recent"] != 0:
                 new = True
         if not new:
@@ -50,7 +50,7 @@ def summary(event, context):
         Destination={"ToAddresses": [environ.get("ADMIN_CONTACT")]},
         Message={
             "Body": {
-                "Html": {"Charset": "UTF-8", "Data": generate(event)},
+                "Html": {"Charset": "UTF-8", "Data": generate()},
                 "Text": {
                     "Charset": "UTF-8",
                     "Data": "This email must be viewed with HTML",
diff --git a/linkshortener/shortener.py b/linkshortener/shortener.py
index 505d6c4..2498e1a 100644
--- a/linkshortener/shortener.py
+++ b/linkshortener/shortener.py
@@ -8,11 +8,11 @@
 from botocore.config import Config
 
 
-def connect(event):
+def connect():
     return boto3.resource(
         "dynamodb",
         endpoint_url="http://localhost:8000"
-        if event["requestContext"]["stage"] == "dev"
+        if os.environ.get("STAGE") == "dev"
         else None,
         region_name="eu-west-2",
         config=Config(connect_timeout=5, retries={"max_attempts": 3}),
@@ -40,7 +40,7 @@ def redirect(url):
 
 
 def shortener(event, context):
-    db = connect(event)
+    db = connect()
     try:
         url = db.get_item(Key={"code": event["pathParameters"]["id"]})["Item"]["url"]
         db.update_item(
@@ -57,7 +57,7 @@ def shortener(event, context):
 
 
 def create(event, context):
-    db = connect(event)
+    db = connect()
     try:
         db.put_item(
             Item={
@@ -74,7 +74,7 @@ def create(event, context):
 
 
 def view(event, context):
-    db = connect(event)
+    db = connect()
     if event["queryStringParameters"] is not None and event[
         "queryStringParameters"
     ].get("code") not in (None, ""):
@@ -90,7 +90,7 @@ def view(event, context):
 
 
 def delete(event, context):
-    db = connect(event)
+    db = connect()
     try:
         db.delete_item(
             Key={"code": sanitize(json.loads(event["body"])["code"])},
@@ -103,7 +103,7 @@ def delete(event, context):
 
 
 def fallback(event, context):
-    return redirect(os.environ["FALLBACK_URL"])
+    return redirect(os.environ.get("FALLBACK_URL"))
 
 
 def robots(event, context):
diff --git a/serverless.yml b/serverless.yml
index 20d96c8..6eb11f3 100644
--- a/serverless.yml
+++ b/serverless.yml
@@ -7,7 +7,7 @@ provider:
   runtime: python3.8
   region: eu-west-2
   memorySize: 128
-  timeout: 20
+  timeout: 6
   apiKeys:
     - ${opt:stage}-adminkey
   environment:
@@ -16,6 +16,7 @@ provider:
     DOMAIN: ${env:DOMAIN}
     SES_REGION: ${env:SES_REGION}
     FALLBACK_URL: ${env:FALLBACK_URL}
+    STAGE: ${opt:stage, self:provider.stage}
   iamRoleStatements:
       - Effect: Allow
         Action: