Skip to content

Commit aa2b54d

Browse files
amanceviceamancevice
committed
Support for API Gateway v2 payloads
1 parent 36f1868 commit aa2b54d

File tree

5 files changed

+319
-65
lines changed

5 files changed

+319
-65
lines changed

‎lambda_gateway/__main__.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,12 @@ def get_opts():
7373
help='Print version and exit',
7474
version=f'%(prog)s {__version__}',
7575
)
76+
parser.add_argument(
77+
'-V', '--payload-version',
78+
choices=['1.0', '2.0'],
79+
default='2.0',
80+
help='API Gateway payload version [default: 2.0]',
81+
)
7682
parser.add_argument(
7783
'HANDLER',
7884
help='Lambda handler signature',
@@ -114,7 +120,7 @@ def main():
114120
# Setup handler
115121
address_family, addr = get_best_family(opts.bind, opts.port)
116122
proxy = EventProxy(opts.HANDLER, base_path, opts.timeout)
117-
LambdaRequestHandler.set_proxy(proxy)
123+
LambdaRequestHandler.set_proxy(proxy, opts.payload_version)
118124
server.ThreadingHTTPServer.address_family = address_family
119125

120126
# Start server

‎lambda_gateway/event_proxy.py

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,28 @@ def get_handler(self):
3333
except AttributeError:
3434
raise ValueError(f"Handler '{func}' missing on module '{name}'")
3535

36+
def get_httpMethod(self, event):
37+
"""
38+
Helper to get httpMethod from v1 or v2 events.
39+
"""
40+
if event.get('version') == '2.0':
41+
return event['requestContext']['http']['method']
42+
elif event.get('version') == '1.0':
43+
return event['httpMethod']
44+
raise ValueError( # pragma: no cover
45+
f"Unknown API Gateway payload version: {event.get('version')}")
46+
47+
def get_path(self, event):
48+
"""
49+
Helper to get path from v1 or v2 events.
50+
"""
51+
if event.get('version') == '2.0':
52+
return event['rawPath']
53+
elif event.get('version') == '1.0':
54+
return event['path']
55+
raise ValueError( # pragma: no cover
56+
f"Unknown API Gateway payload version: {event.get('version')}")
57+
3658
def invoke(self, event):
3759
with lambda_context.start(self.timeout) as context:
3860
logger.info('Invoking "%s"', self.handler)
@@ -46,8 +68,8 @@ async def invoke_async(self, event, context=None):
4668
:param Context context: Mock Lambda context
4769
:returns dict: Lamnda invocation result
4870
"""
49-
httpMethod = event['httpMethod']
50-
path = event['path']
71+
httpMethod = self.get_httpMethod(event)
72+
path = self.get_path(event)
5173

5274
# Reject request if not starting at base path
5375
if not path.startswith(self.base_path):
@@ -77,7 +99,7 @@ async def invoke_async_with_timeout(self, event, context=None):
7799
coroutine = self.invoke_async(event, context)
78100
return await asyncio.wait_for(coroutine, self.timeout)
79101
except asyncio.TimeoutError:
80-
httpMethod = event['httpMethod']
102+
httpMethod = self.get_httpMethod(event)
81103
message = 'Endpoint request timed out'
82104
return self.jsonify(httpMethod, 504, message=message)
83105

‎lambda_gateway/request_handler.py

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,31 @@ def get_event(self, httpMethod):
2626
"""
2727
Get Lambda input event object.
2828
29+
:param str httpMethod: HTTP request method
30+
:return dict: Lambda event object
31+
"""
32+
if self.version == '1.0':
33+
return self.get_event_v1(httpMethod)
34+
elif self.version == '2.0':
35+
return self.get_event_v2(httpMethod)
36+
raise ValueError( # pragma: no cover
37+
f'Unknown API Gateway payload version: {self.version}')
38+
39+
def get_event_v1(self, httpMethod):
40+
"""
41+
Get Lambda input event object (v1).
42+
2943
:param str httpMethod: HTTP request method
3044
:return dict: Lambda event object
3145
"""
3246
url = parse.urlparse(self.path)
47+
path, *_ = url.path.split('?')
3348
return {
49+
'version': '1.0',
3450
'body': self.get_body(),
3551
'headers': dict(self.headers),
3652
'httpMethod': httpMethod,
37-
'path': url.path,
53+
'path': path,
3854
'queryStringParameters': dict(parse.parse_qsl(url.query)),
3955
}
4056

@@ -46,18 +62,19 @@ def get_event_v2(self, httpMethod):
4662
:return dict: Lambda event object
4763
"""
4864
url = parse.urlparse(self.path)
65+
path, *_ = url.path.split('?')
4966
return {
5067
'version': '2.0',
5168
'body': self.get_body(),
52-
'routeKey': f'{httpMethod} {url.path}',
53-
'rawPath': url.path,
69+
'routeKey': f'{httpMethod} {path}',
70+
'rawPath': path,
5471
'rawQueryString': url.query,
5572
'headers': dict(self.headers),
5673
'queryStringParameters': dict(parse.parse_qsl(url.query)),
5774
'requestContext': {
5875
'http': {
5976
'method': httpMethod,
60-
'path': url.path,
77+
'path': path,
6178
},
6279
},
6380
}
@@ -89,8 +106,9 @@ def invoke(self, httpMethod):
89106
self.wfile.write(body.encode())
90107

91108
@classmethod
92-
def set_proxy(cls, proxy):
109+
def set_proxy(cls, proxy, version):
93110
"""
94111
Set up LambdaRequestHandler.
95112
"""
96113
cls.proxy = proxy
114+
cls.version = version

‎tests/test_event_proxy.py

Lines changed: 198 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -28,44 +28,218 @@ def test_get_handler_error(self, handler):
2828
with pytest.raises(ValueError):
2929
self.subject.get_handler()
3030

31-
@pytest.mark.parametrize(('verb', 'path', 'status', 'message'), [
32-
('GET', '/simple/', 200, 'OK'),
33-
('HEAD', '/simple/', 200, 'OK'),
34-
('POST', '/simple/', 200, 'OK'),
31+
@pytest.mark.parametrize(('event', 'exp'), [
32+
(
33+
{
34+
'version': '1.0',
35+
'httpMethod': 'GET',
36+
'path': '/simple/',
37+
},
38+
EventProxy.jsonify('GET', 200, message='OK'),
39+
),
40+
(
41+
{
42+
'version': '1.0',
43+
'httpMethod': 'HEAD',
44+
'path': '/simple/',
45+
},
46+
EventProxy.jsonify('HEAD', 200, message='OK'),
47+
),
48+
(
49+
{
50+
'version': '1.0',
51+
'httpMethod': 'POST',
52+
'path': '/simple/',
53+
'body': '{"fizz": "buzz"}',
54+
},
55+
EventProxy.jsonify('POST', 200, message='OK'),
56+
),
57+
(
58+
{
59+
'version': '2.0',
60+
'rawPath': '/simple/',
61+
'requestContext': {
62+
'http': {
63+
'method': 'GET',
64+
'path': '/simple/',
65+
},
66+
},
67+
},
68+
EventProxy.jsonify('GET', 200, message='OK'),
69+
),
70+
(
71+
{
72+
'version': '2.0',
73+
'rawPath': '/simple/',
74+
'requestContext': {
75+
'http': {
76+
'method': 'HEAD',
77+
'path': '/simple/',
78+
},
79+
},
80+
},
81+
EventProxy.jsonify('HEAD', 200, message='OK'),
82+
),
83+
(
84+
{
85+
'version': '2.0',
86+
'rawPath': '/simple/',
87+
'body': '{"fizz": "buzz"}',
88+
'requestContext': {
89+
'http': {
90+
'method': 'POST',
91+
'path': '/simple/',
92+
},
93+
},
94+
},
95+
EventProxy.jsonify('POST', 200, message='OK'),
96+
),
3597
])
36-
def test_invoke_success(self, verb, path, status, message):
98+
def test_invoke_success(self, event, exp):
3799
self.subject.get_handler = lambda: lambda event, context: exp
38-
exp = EventProxy.jsonify(verb, status, message=message)
39-
ret = self.subject.invoke({'httpMethod': verb, 'path': path})
100+
ret = self.subject.invoke(event)
40101
assert ret == exp
41102

42-
@pytest.mark.parametrize(('verb', 'path', 'status', 'message'), [
43-
('GET', '/simple/', 502, 'Internal server error'),
44-
('HEAD', '/simple/', 502, 'Internal server error'),
45-
('POST', '/simple/', 502, 'Internal server error'),
46-
('GET', '/', 403, 'Forbidden'),
47-
('HEAD', '/', 403, 'Forbidden'),
48-
('POST', '/', 403, 'Forbidden'),
103+
@pytest.mark.parametrize(('event', 'exp'), [
104+
(
105+
{
106+
'version': '2.0',
107+
'rawPath': '/simple/',
108+
'requestContext': {
109+
'http': {
110+
'method': 'GET',
111+
'path': '/simple/',
112+
},
113+
},
114+
},
115+
EventProxy.jsonify('GET', 502, message='Internal server error'),
116+
),
117+
(
118+
{
119+
'version': '2.0',
120+
'rawPath': '/simple/',
121+
'requestContext': {
122+
'http': {
123+
'method': 'HEAD',
124+
'path': '/simple/',
125+
},
126+
},
127+
},
128+
EventProxy.jsonify('HEAD', 502, message='Internal server error'),
129+
),
130+
(
131+
{
132+
'version': '2.0',
133+
'rawPath': '/simple/',
134+
'requestContext': {
135+
'http': {
136+
'method': 'POST',
137+
'path': '/simple/',
138+
},
139+
},
140+
},
141+
EventProxy.jsonify('POST', 502, message='Internal server error'),
142+
),
143+
(
144+
{
145+
'version': '2.0',
146+
'rawPath': '/',
147+
'requestContext': {
148+
'http': {
149+
'method': 'GET',
150+
'path': '/',
151+
},
152+
},
153+
},
154+
EventProxy.jsonify('GET', 403, message='Forbidden'),
155+
),
156+
(
157+
{
158+
'version': '2.0',
159+
'rawPath': '/',
160+
'requestContext': {
161+
'http': {
162+
'method': 'HEAD',
163+
'path': '/',
164+
},
165+
},
166+
},
167+
EventProxy.jsonify('HEAD', 403, message='Forbidden'),
168+
),
169+
(
170+
{
171+
'version': '2.0',
172+
'rawPath': '/',
173+
'requestContext': {
174+
'http': {
175+
'method': 'POST',
176+
'path': '/',
177+
},
178+
},
179+
},
180+
EventProxy.jsonify('POST', 403, message='Forbidden'),
181+
),
49182
])
50-
def test_invoke_error(self, verb, path, status, message):
183+
def test_invoke_error(self, event, exp):
51184
def handler(event, context):
52185
raise Exception()
53186
self.subject.get_handler = lambda: handler
54-
exp = EventProxy.jsonify(verb, status, message=message)
55-
ret = self.subject.invoke({'httpMethod': verb, 'path': path})
187+
ret = self.subject.invoke(event)
56188
assert ret == exp
57189

58-
@pytest.mark.parametrize(('verb', 'path', 'status', 'message'), [
59-
('GET', '/simple/', 504, 'Endpoint request timed out'),
60-
('HEAD', '/simple/', 504, 'Endpoint request timed out'),
61-
('POST', '/simple/', 504, 'Endpoint request timed out'),
190+
@pytest.mark.parametrize(('event', 'exp'), [
191+
(
192+
{
193+
'version': '2.0',
194+
'rawPath': '/simple/',
195+
'requestContext': {
196+
'http': {
197+
'method': 'GET',
198+
'path': '/simple/',
199+
},
200+
},
201+
},
202+
EventProxy.jsonify(
203+
'GET', 504,
204+
message='Endpoint request timed out',
205+
),
206+
),
207+
(
208+
{
209+
'version': '2.0',
210+
'rawPath': '/simple/',
211+
'requestContext': {
212+
'http': {
213+
'method': 'HEAD',
214+
'path': '/simple/',
215+
},
216+
},
217+
},
218+
EventProxy.jsonify('HEAD', 504),
219+
),
220+
(
221+
{
222+
'version': '2.0',
223+
'rawPath': '/simple/',
224+
'body': '{"fizz": "buzz"}',
225+
'requestContext': {
226+
'http': {
227+
'method': 'POST',
228+
'path': '/simple/',
229+
},
230+
},
231+
},
232+
EventProxy.jsonify(
233+
'POST', 504,
234+
message='Endpoint request timed out',
235+
),
236+
),
62237
])
63-
def test_invoke_timeout(self, verb, path, status, message):
238+
def test_invoke_timeout(self, event, exp):
64239
patch = 'lambda_gateway.event_proxy.EventProxy.invoke_async'
65240
with mock.patch(patch) as mock_invoke:
66241
mock_invoke.side_effect = asyncio.TimeoutError
67-
exp = EventProxy.jsonify(verb, status, message=message)
68-
ret = self.subject.invoke({'httpMethod': verb, 'path': path})
242+
ret = self.subject.invoke(event)
69243
assert ret == exp
70244

71245
@pytest.mark.parametrize(('verb', 'statusCode', 'body', 'exp'), [

0 commit comments

Comments
 (0)