add tests and wrappers

Author: Joe Snell

bin/format

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+set -e
+
+poetry run black .

bin/test

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+set -e
+
+poetry run pytest

lef/__init__.py

@@ -5,47 +5,46 @@
 from pathlib import Path
 from enum import Enum
 
-from lambda_extension_framework.utils import get, post
+from lef.utils import get, post
+
 
 class EventType(Enum):
-    INVOKE='INVOKE'
-    SHUTDOWN='SHUTDOWN'
+    INVOKE = "INVOKE"
+    SHUTDOWN = "SHUTDOWN"
+
 
 class Extension:
     def __init__(self, name=None, runtime_api_endpoint=None):
         self.handlers = {}
         self.name = name or Path(__file__).parent.name
-        self.runtime_api_endpoint = runtime_api_endpoint or os.environ['AWS_LAMBDA_RUNTIME_API']
+        self.runtime_api_endpoint = (
+            runtime_api_endpoint or os.environ["AWS_LAMBDA_RUNTIME_API"]
+        )
 
     @staticmethod
     def handle_signal(signal, frame):
         sys.exit(0)
 
     def register(self, events, handler):
         response = post(
-            url=f'https://{self.runtime_api_endpoint}/2020-01-01/extension/register',
-            data={
-                'events': events
-            },
+            url=f"https://{self.runtime_api_endpoint}/2020-01-01/extension/register",
+            data={"events": [str(event) for event in events]},
             headers={
-                'Lambda-Extension-Name': self.name,
-            }
+                "Lambda-Extension-Name": self.name,
+            },
         )
-    
-        extension_id = response.headers['Lambda-Extension-Identifier']
+
+        extension_id = response.headers["Lambda-Extension-Identifier"]
         return self.process_events(extension_id, handler)
 
     def process_events(self, extension_id, handler):
         while True:
             response = get(
                 url=f"https://{self.runtime_api_endpoint}/2020-01-01/extension/event/next",
-                headers={
-                    'Lambda-Extension-Identifier': extension_id
-                },
-                timeout=None
-            )
+                headers={"Lambda-Extension-Identifier": extension_id},
+            ).json()
 
-            if response['eventType'] == EventType.SHUTDOWN:
+            if response["eventType"] == EventType.SHUTDOWN:
                 sys.exit(0)
             else:
                 handler(response)

lef/main.py

@@ -1,26 +1,31 @@
 import json
 from urllib import request
 
+
+class Response:
+    def __init__(self, res):
+        self.headers = res.headers
+        self.text = res.read().decode()
+
+    def json(self):
+        return json.loads(self.text)
+
+
 def get(url, headers=None):
-    kwargs = {
-        'url': url,
-        'headers': headers or {}
-    }
+    kwargs = {"url": url, "headers": headers or {}}
 
-    req = request.Request(**kwargs)                    
+    req = request.Request(**kwargs)
     response = request.urlopen(req)
-    return response.read().decode('utf8')
+    return Response(response)
+
 
 def post(url, headers=None, data=None):
-    kwargs = {
-        'url': url,
-        'headers': headers or {}
-    }
+    kwargs = {"url": url, "headers": headers or {}}
 
     if data:
-        kwargs['data'] = json.dumps(data)
-        kwargs['headers']['Content-Type'] = 'application/json'
+        kwargs["data"] = json.dumps(data)
+        kwargs["headers"]["Content-Type"] = "application/json"
 
-    req = request.Request(**kwargs)                    
+    req = request.Request(**kwargs)
     response = request.urlopen(req)
-    return response.read().decode('utf8')
+    return Response(response)