Skip to content

Commit e1cf67a

Browse files
nathangloverkhamaileon
nathanglover
authored and
khamaileon
committed
Kludex#181 - ensure bytes are always returned from .body() method (Kludex#182)
* Kludex#181 - ensure bytes are always returned from .body method * Kludex#181 - remove extra if * Kludex#181 - black formatting * Kludex#181 - remove extra if
1 parent 78432b5 commit e1cf67a

File tree

8 files changed

+80
-16
lines changed

8 files changed

+80
-16
lines changed

‎mangum/handlers/aws_alb.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,9 +57,13 @@ def request(self) -> Request:
5757

5858
@property
5959
def body(self) -> bytes:
60-
body = self.trigger_event.get("body", b"")
60+
body = self.trigger_event.get("body", b"") or b""
61+
6162
if self.trigger_event.get("isBase64Encoded", False):
62-
body = base64.b64decode(body)
63+
return base64.b64decode(body)
64+
if not isinstance(body, bytes):
65+
body = body.encode()
66+
6367
return body
6468

6569
def transform_response(self, response: Response) -> Dict[str, Any]:

‎mangum/handlers/aws_api_gateway.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,9 +95,13 @@ def request(self) -> Request:
9595

9696
@property
9797
def body(self) -> bytes:
98-
body = self.trigger_event.get("body", b"")
98+
body = self.trigger_event.get("body", b"") or b""
99+
99100
if self.trigger_event.get("isBase64Encoded", False):
100-
body = base64.b64decode(body)
101+
return base64.b64decode(body)
102+
if not isinstance(body, bytes):
103+
body = body.encode()
104+
101105
return body
102106

103107
def transform_response(self, response: Response) -> Dict[str, Any]:

‎mangum/handlers/aws_cf_lambda_at_edge.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,13 @@ def request(self) -> Request:
5555
@property
5656
def body(self) -> bytes:
5757
request = self.trigger_event["Records"][0]["cf"]["request"]
58-
body = request.get("body", {}).get("data", None)
58+
body = request.get("body", {}).get("data", None) or b""
59+
5960
if request.get("body", {}).get("encoding", "") == "base64":
60-
body = base64.b64decode(body)
61+
return base64.b64decode(body)
62+
if not isinstance(body, bytes):
63+
body = body.encode()
64+
6165
return body
6266

6367
def transform_response(self, response: Response) -> Dict[str, Any]:

‎mangum/handlers/aws_http_gateway.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,13 @@ def request(self) -> Request:
102102

103103
@property
104104
def body(self) -> bytes:
105-
body = self.trigger_event.get("body", b"")
105+
body = self.trigger_event.get("body", b"") or b""
106+
106107
if self.trigger_event.get("isBase64Encoded", False):
107-
body = base64.b64decode(body)
108+
return base64.b64decode(body)
109+
if not isinstance(body, bytes):
110+
body = body.encode()
111+
108112
return body
109113

110114
def transform_response(self, response: Response) -> Dict[str, Any]:

‎tests/handlers/test_aws_alb.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,8 @@ def test_aws_alb_basic():
7474

7575
example_context = {}
7676
handler = AwsAlb(example_event, example_context)
77+
78+
assert type(handler.body) == bytes
7779
assert handler.request.scope == {
7880
"asgi": {"version": "3.0"},
7981
"aws.context": {},
@@ -119,7 +121,15 @@ def test_aws_alb_basic():
119121
"query_string,scope_body",
120122
[
121123
("GET", "/hello/world", None, None, False, b"", None),
122-
("POST", "/", {"name": ["me"]}, None, False, b"name=me", None),
124+
(
125+
"POST",
126+
"/",
127+
{"name": ["me"]},
128+
"field1=value1&field2=value2",
129+
False,
130+
b"name=me",
131+
b"field1=value1&field2=value2",
132+
),
123133
(
124134
"GET",
125135
"/my/resource",
@@ -210,7 +220,10 @@ def test_aws_alb_scope_real(
210220
"type": "http",
211221
}
212222

213-
assert handler.body == scope_body
223+
if handler.body:
224+
assert handler.body == scope_body
225+
else:
226+
assert handler.body == b""
214227

215228

216229
def test_aws_alb_set_cookies() -> None:

‎tests/handlers/test_aws_api_gateway.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,8 @@ def test_aws_api_gateway_scope_basic():
9696
}
9797
example_context = {}
9898
handler = AwsApiGateway(example_event, example_context)
99+
100+
assert type(handler.body) == bytes
99101
assert handler.request.scope == {
100102
"asgi": {"version": "3.0"},
101103
"aws.context": {},
@@ -128,7 +130,15 @@ def test_aws_api_gateway_scope_basic():
128130
"query_string,scope_body",
129131
[
130132
("GET", "/hello/world", None, None, False, b"", None),
131-
("POST", "/", {"name": ["me"]}, None, False, b"name=me", None),
133+
(
134+
"POST",
135+
"/",
136+
{"name": ["me"]},
137+
"field1=value1&field2=value2",
138+
False,
139+
b"name=me",
140+
b"field1=value1&field2=value2",
141+
),
132142
(
133143
"GET",
134144
"/my/resource",
@@ -218,7 +228,10 @@ def test_aws_api_gateway_scope_real(
218228
"type": "http",
219229
}
220230

221-
assert handler.body == scope_body
231+
if handler.body:
232+
assert handler.body == scope_body
233+
else:
234+
assert handler.body == b""
222235

223236

224237
@pytest.mark.parametrize(

‎tests/handlers/test_aws_cf_lambda_at_edge.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ def test_aws_cf_lambda_at_edge_scope_basic():
136136
example_context = {}
137137
handler = AwsCfLambdaAtEdge(example_event, example_context)
138138

139+
assert type(handler.body) == bytes
139140
assert handler.request.scope == {
140141
"asgi": {"version": "3.0"},
141142
"aws.context": {},
@@ -169,7 +170,15 @@ def test_aws_cf_lambda_at_edge_scope_basic():
169170
"body_base64_encoded,query_string,scope_body",
170171
[
171172
("GET", "/hello/world", None, None, False, b"", None),
172-
("POST", "/", {"name": ["me"]}, None, False, b"name=me", None),
173+
(
174+
"POST",
175+
"/",
176+
{"name": ["me"]},
177+
"field1=value1&field2=value2",
178+
False,
179+
b"name=me",
180+
b"field1=value1&field2=value2",
181+
),
173182
(
174183
"GET",
175184
"/my/resource",
@@ -241,7 +250,10 @@ def test_aws_api_gateway_scope_real(
241250
"type": "http",
242251
}
243252

244-
assert handler.body == scope_body
253+
if handler.body:
254+
assert handler.body == scope_body
255+
else:
256+
assert handler.body == b""
245257

246258

247259
@pytest.mark.parametrize(

‎tests/handlers/test_aws_http_gateway.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,8 @@ def test_aws_http_gateway_scope_basic_v1():
195195
}
196196
example_context = {}
197197
handler = AwsHttpGateway(example_event, example_context)
198+
199+
assert type(handler.body) == bytes
198200
assert handler.request.scope == {
199201
"asgi": {"version": "3.0"},
200202
"aws.context": {},
@@ -297,6 +299,8 @@ def test_aws_http_gateway_scope_basic_v2():
297299
}
298300
example_context = {}
299301
handler = AwsHttpGateway(example_event, example_context)
302+
303+
assert type(handler.body) == bytes
300304
assert handler.request.scope == {
301305
"asgi": {"version": "3.0"},
302306
"aws.context": {},
@@ -396,7 +400,10 @@ def test_aws_http_gateway_scope_real_v1(
396400
"type": "http",
397401
}
398402

399-
assert handler.body == scope_body
403+
if handler.body:
404+
assert handler.body == scope_body
405+
else:
406+
assert handler.body == b""
400407

401408

402409
@pytest.mark.parametrize(
@@ -461,7 +468,10 @@ def test_aws_http_gateway_scope_real_v2(
461468
"type": "http",
462469
}
463470

464-
assert handler.body == scope_body
471+
if handler.body:
472+
assert handler.body == scope_body
473+
else:
474+
assert handler.body == b""
465475

466476

467477
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)
⚡