diff --git a/src/granule_to_sns.py b/src/granule_to_sns.py
index f30e297..cc7caab 100644
--- a/src/granule_to_sns.py
+++ b/src/granule_to_sns.py
@@ -27,9 +27,11 @@ def generate_message(granule: dict) -> dict:
 
 
 def granule_to_sns(event: dict, _) -> dict:
-    client = boto3.client("sns")
-    granules = event["input"]["granules"]
+    payload = event["input"]
+
+    granules = payload["granules"]
 
+    client = boto3.client("sns")
     sns_topic_arn = os.getenv("SNS_TOPIC_ARN")
 
     for granule in granules:
@@ -46,10 +48,10 @@ def granule_to_sns(event: dict, _) -> dict:
             },
         )
 
-    return event
+    return payload
 
 
-def lambda_handler(event, context):
+def lambda_handler(event: dict, context) -> dict:
     init_root_logger()
     with log_errors():
         return run_cumulus_task(granule_to_sns, event, context)
diff --git a/tests/test_granule_to_sns.py b/tests/test_granule_to_sns.py
index ccd1d35..618ddb2 100644
--- a/tests/test_granule_to_sns.py
+++ b/tests/test_granule_to_sns.py
@@ -1,6 +1,5 @@
 import boto3
 import pytest
-
 from granule_to_sns import generate_message, granule_to_sns
 
 
@@ -105,4 +104,4 @@ def test_generate_message(event, message):
 
 def test_granule_to_sns(sns_client, event, mocker):
     mocker.patch.object(sns_client, "publish", return_value={})
-    assert granule_to_sns(event, None) == event
+    assert granule_to_sns(event, None) == event["input"]