get tests working

Author: Joe Snell

bin/test

@@ -2,4 +2,4 @@
 
 set -e
 
-poetry run pytest
+poetry run pytest "$@"

lef/main.py

@@ -26,6 +26,7 @@ def handle_signal(signal, frame):
         sys.exit(0)
 
     def register(self, events, handler):
+        print('called')
         response = post(
             url=f"https://{self.runtime_api_endpoint}/2020-01-01/extension/register",
             data={"events": [str(event) for event in events]},
@@ -39,12 +40,13 @@ def register(self, events, handler):
 
     def process_events(self, extension_id, handler):
         while True:
+            print('called')
             response = get(
                 url=f"https://{self.runtime_api_endpoint}/2020-01-01/extension/event/next",
                 headers={"Lambda-Extension-Identifier": extension_id},
             ).json()
 
-            if response["eventType"] == EventType.SHUTDOWN:
+            if response["eventType"] == 'SHUTDOWN':
                 sys.exit(0)
             else:
                 handler(response)

tests/test_main.py

@@ -2,9 +2,9 @@
 
 from lef.main import EventType, Extension
 from unittest import TestCase
-from unittest.mock import patch, call
+from unittest.mock import MagicMock, patch, call
 from tests.utils import FakeResponse
-
+import pytest
 
 class MainTest(TestCase):
     def setUp(self):
@@ -15,19 +15,39 @@ def tearDown(self):
         os.environ = self.original_env
         self.original_env = None
 
-    @patch("lef.utils.request")
-    def test_main(self, mock_request):
-        mock_request.Request.return_value = "FakeRequestObject"
-        mock_request.urlopen.return_value = FakeResponse(
-            json.dumps({"eventType": "INVOKE"}).encode(),
-            {"Lambda-Extension-Identifier": "lambda-ext-id"},
-        )
-
-        def handler(event):
-            assert event == {"eventType": "INVOKE"}
-
+    @patch('lef.main.post')
+    @patch("lef.main.get")
+    def test_main(self, mock_get, mock_post):
+        mock_get.side_effect = [
+            FakeResponse(
+                json.dumps({"eventType": "INVOKE"})
+            ),
+            FakeResponse(
+                json.dumps({"eventType": "SHUTDOWN"})
+            )
+        ]
+
+        mock_post.side_effect = [
+            FakeResponse(
+                '',
+                {"Lambda-Extension-Identifier": "lambda-ext-id"}
+            ),
+        ]
+
+        handler = MagicMock()
         extension = Extension()
-        extension.register([EventType.INVOKE], handler)
+
+        with pytest.raises(SystemExit) as e:
+            extension.register([EventType.INVOKE], handler)
 
         assert extension.runtime_api_endpoint == "testruntimeapi"
-        assert [] == mock_request.mock_calls
+        assert [
+            call({'eventType': 'INVOKE'})
+        ] == handler.mock_calls
+        assert [
+            call(url='https://testruntimeapi/2020-01-01/extension/event/next', headers={'Lambda-Extension-Identifier': 'lambda-ext-id'}),
+            call(url='https://testruntimeapi/2020-01-01/extension/event/next', headers={'Lambda-Extension-Identifier': 'lambda-ext-id'})
+        ] == mock_get.mock_calls
+        assert [
+            call(url='https://testruntimeapi/2020-01-01/extension/register', data={'events': ['EventType.INVOKE']}, headers={'Lambda-Extension-Name': 'lef'})
+        ] == mock_post.mock_calls

tests/test_utils.py

@@ -5,6 +5,7 @@
 
 
 class UtilsTest(TestCase):
+
     @patch("lef.utils.request")
     def test_get(self, mock_request):
         mock_request.Request.return_value = "FakeRequestObject"

tests/utils.py

@@ -1,3 +1,5 @@
+import json
+
 class FakeResponse:
     def __init__(self, response, headers=None):
         self._response = response
@@ -6,3 +8,7 @@ def __init__(self, response, headers=None):
 
     def read(self):
         return self._response
+
+    def json(self):
+        return json.loads(self._response)
+