Create a new set_body function and use it everywhere

aikido_zen/context/__init__.py

@@ -3,6 +3,7 @@
 """
 
 import contextvars
+import json
 from urllib.parse import parse_qs
 
 from aikido_zen.helpers.build_route_from_url import build_route_from_url
@@ -44,7 +45,7 @@ def __init__(self, context_obj=None, body=None, req=None, source=None):
         self.parsed_userinput = {}
         self.xml = {}
         self.outgoing_req_redirects = []
-        self.body = body
+        self.set_body(body)
 
         # Parse WSGI/ASGI/... request :
         self.cookies = self.method = self.remote_address = self.query = self.headers = (
@@ -91,6 +92,14 @@ def set_as_current_context(self):
         """
         current_context.set(self)
 
+    def set_body(self, body):
+        """Sets the body, and verifies the body is okay"""
+        try:
+            json.dumps(body)
+            self.body = body
+        except (TypeError, OverflowError):
+            self.body = None
+
     def get_route_metadata(self):
         """Returns a route_metadata object"""
         return {

aikido_zen/context/init_test.py

@@ -4,6 +4,21 @@
 from aikido_zen.context import Context, get_current_context, current_context
 
 
+basic_wsgi_req = {
+    "REQUEST_METHOD": "GET",
+    "HTTP_HEADER_1": "header 1 value",
+    "HTTP_HEADER_2": "Header 2 value",
+    "RANDOM_VALUE": "Random value",
+    "HTTP_COOKIE": "sessionId=abc123xyz456;",
+    "wsgi.url_scheme": "http",
+    "HTTP_HOST": "localhost:8080",
+    "PATH_INFO": "/hello",
+    "QUERY_STRING": "user=JohnDoe&age=30&age=35",
+    "CONTENT_TYPE": "application/json",
+    "REMOTE_ADDR": "198.51.100.23",
+}
+
+
 @pytest.fixture(autouse=True)
 def run_around_tests():
     yield
@@ -58,20 +73,7 @@ def test_wsgi_context_1():
 
 
 def test_wsgi_context_2():
-    wsgi_request = {
-        "REQUEST_METHOD": "GET",
-        "HTTP_HEADER_1": "header 1 value",
-        "HTTP_HEADER_2": "Header 2 value",
-        "RANDOM_VALUE": "Random value",
-        "HTTP_COOKIE": "sessionId=abc123xyz456;",
-        "wsgi.url_scheme": "http",
-        "HTTP_HOST": "localhost:8080",
-        "PATH_INFO": "/hello",
-        "QUERY_STRING": "user=JohnDoe&age=30&age=35",
-        "CONTENT_TYPE": "application/json",
-        "REMOTE_ADDR": "198.51.100.23",
-    }
-    context = Context(req=wsgi_request, body={"test": True}, source="flask")
+    context = Context(req=basic_wsgi_req, body={"test": True}, source="flask")
     assert context.__dict__ == {
         "source": "flask",
         "method": "GET",
@@ -99,59 +101,20 @@ def test_wsgi_context_2():
 
 def test_set_as_current_context(mocker):
     # Test set_as_current_context() method
-    wsgi_request = {
-        "REQUEST_METHOD": "GET",
-        "HTTP_HEADER_1": "header 1 value",
-        "HTTP_HEADER_2": "Header 2 value",
-        "RANDOM_VALUE": "Random value",
-        "HTTP_COOKIE": "sessionId=abc123xyz456;",
-        "wsgi.url_scheme": "http",
-        "HTTP_HOST": "localhost:8080",
-        "PATH_INFO": "/hello",
-        "QUERY_STRING": "user=JohnDoe&age=30&age=35",
-        "CONTENT_TYPE": "application/json",
-        "REMOTE_ADDR": "198.51.100.23",
-    }
-    context = Context(req=wsgi_request, body=12, source="flask")
+    context = Context(req=basic_wsgi_req, body=12, source="flask")
     context.set_as_current_context()
     assert get_current_context() == context
 
 
 def test_get_current_context_with_context(mocker):
     # Test get_current_context() when a context is set
-    wsgi_request = {
-        "REQUEST_METHOD": "GET",
-        "HTTP_HEADER_1": "header 1 value",
-        "HTTP_HEADER_2": "Header 2 value",
-        "RANDOM_VALUE": "Random value",
-        "HTTP_COOKIE": "sessionId=abc123xyz456;",
-        "wsgi.url_scheme": "http",
-        "HTTP_HOST": "localhost:8080",
-        "PATH_INFO": "/hello",
-        "QUERY_STRING": "user=JohnDoe&age=30&age=35",
-        "CONTENT_TYPE": "application/json",
-        "REMOTE_ADDR": "198.51.100.23",
-    }
-    context = Context(req=wsgi_request, body=456, source="flask")
+    context = Context(req=basic_wsgi_req, body=456, source="flask")
     context.set_as_current_context()
     assert get_current_context() == context
 
 
 def test_context_is_picklable(mocker):
-    wsgi_request = {
-        "REQUEST_METHOD": "GET",
-        "HTTP_HEADER_1": "header 1 value",
-        "HTTP_HEADER_2": "Header 2 value",
-        "RANDOM_VALUE": "Random value",
-        "HTTP_COOKIE": "sessionId=abc123xyz456;",
-        "wsgi.url_scheme": "http",
-        "HTTP_HOST": "localhost:8080",
-        "PATH_INFO": "/hello",
-        "QUERY_STRING": "user=JohnDoe&age=30&age=35",
-        "CONTENT_TYPE": "application/json",
-        "REMOTE_ADDR": "198.51.100.23",
-    }
-    context = Context(req=wsgi_request, body=123, source="flask")
+    context = Context(req=basic_wsgi_req, body=123, source="flask")
     pickled_obj = pickle.dumps(context)
     unpickled_obj = pickle.loads(pickled_obj)
     assert unpickled_obj.source == "flask"
@@ -168,3 +131,38 @@ def test_context_is_picklable(mocker):
     }
     assert unpickled_obj.query == {"user": ["JohnDoe"], "age": ["30", "35"]}
     assert unpickled_obj.cookies == {"sessionId": "abc123xyz456"}
+
+
+def test_set_valid_dict():
+    valid_body = {"key": "value"}
+    context = Context(req=basic_wsgi_req, body=valid_body, source="flask")
+    assert context.body == valid_body
+
+
+def test_set_valid_list():
+    valid_body = [1, 2, 3]
+    context = Context(req=basic_wsgi_req, body=valid_body, source="flask")
+    assert context.body == valid_body
+
+
+def test_set_valid_string():
+    valid_body = "This is a valid string"
+    context = Context(req=basic_wsgi_req, body=valid_body, source="flask")
+    assert context.body == valid_body
+
+
+def test_set_invalid_body():
+    invalid_body = set([1, 2, 3])  # Sets are not JSON serializable
+    context = Context(req=basic_wsgi_req, body=invalid_body, source="flask")
+    assert context.body is None
+
+
+def test_set_bytestring():
+    invalid_body = b"Hello World"  # Byte strings are not JSON Serializablec
+    context = Context(req=basic_wsgi_req, body=invalid_body, source="flask")
+    assert context.body is None
+
+
+def test_set_none():
+    context = Context(req=basic_wsgi_req, body=None, source="flask")
+    assert context.body is None

aikido_zen/sources/flask.py

@@ -45,11 +45,11 @@ def extract_and_save_data_from_flask_request(req):
         context = get_current_context()
         if context:
             if req.is_json:
-                context.body = req.get_json()
+                context.set_body(req.get_json())
             elif req.form:
-                context.body = req.form
+                context.set_body(req.form)
             else:
-                context.body = req.data.decode("utf-8")
+                context.set_body(req.data.decode("utf-8"))
 
             if getattr(req, "view_args"):
                 context.route_params = dict(req.view_args)

aikido_zen/sources/quart.py

@@ -38,12 +38,12 @@ async def handle_request_wrapper(former_handle_request, quart_app, req):
         if context:
             form = await req.form
             if req.is_json:
-                context.body = await req.get_json()
+                context.set_body(await req.get_json())
             elif form:
-                context.body = form
+                context.set_body(form)
             else:
                 data = await req.data
-                context.body = data.decode("utf-8")
+                context.set_body(data.decode("utf-8"))
             context.cookies = req.cookies.to_dict()
             context.set_as_current_context()
     except Exception as e:

aikido_zen/sources/starlette/extract_data_from_request.py

@@ -11,16 +11,16 @@ async def extract_data_from_request(request):
 
     # Parse data
     try:
-        context.body = await request.json()
+        context.set_body(await request.json())
     except ValueError:
         # Throws error if the body is not json
         pass
     if not context.body:
         form_data = await request.form()
         if form_data:
             # Convert to dict object :
-            context.body = {key: value for key, value in form_data.items()}
+            context.set_body({key: value for key, value in form_data.items()})
     if not context.body:
-        context.body = await request.body()
+        context.set_body(await request.body())
 
     context.set_as_current_context()