Move tests to test_evaluation_plane_handler.

Author: jakeichikawasalesforce

tests/integration/test_max_request_size.py

@@ -60,8 +60,7 @@ def _get_config_file_name(self) -> str:
         return config_file.name
 
     def create_large_payload(self):
-        size_mb = 2
-        num_chars = size_mb * 1024 * 1024
+        num_chars = 2 * 1024 * 1024 # 2MB Size
         large_string = string.printable * (num_chars // len(string.printable))
         large_string += string.printable[:num_chars % len(string.printable)]
         payload = {
@@ -70,26 +69,10 @@ def create_large_payload(self):
         }
         return json.dumps(payload).encode('utf-8')
 
-    def test_payload_exceeds_max_request_size_evaluate(self):
-        headers = { "Content-Type": "application/json" }
-        url = self._get_url() + "/evaluate"
-        response = requests.post(url, data=self.create_large_payload(), headers=headers)
-        self.assertEqual(413, response.status_code)
-
     def test_payload_exceeds_max_request_size_query(self):
         headers = { "Content-Type": "application/json" }
         url = self._get_url() + "/query/model_name"
         response = requests.post(url, data=self.create_large_payload(), headers=headers)
         self.assertEqual(413, response.status_code)
         response = requests.get(url, data=self.create_large_payload(), headers=headers)
         self.assertEqual(413, response.status_code)
-
-    def test_no_content_length_header_present(self):
-        headers = { "Content-Type": "application/json" }
-        url = self._get_url() + "/evaluate"
-        response = requests.post(url, headers=headers)
-        message = json.loads(response.text)["message"]
-        # Ensure it reaches script processing stage in EvaluationPlaneHandler.post
-        self.assertEqual("Error processing script", message)
-        self.assertEqual(500, response.status_code)
-

tests/unit/server_tests/test_evaluation_plane_handler.py

@@ -1,6 +1,8 @@
 import base64
+import json
 import os
 import tempfile
+import string
 
 from tornado.testing import AsyncHTTPTestCase
 
@@ -458,6 +460,67 @@ def test_evaluation_enabled(self):
         )
         self.assertEqual(200, response.code)
 
+class TestEvaluationPlaneHandlerMaxRequestSize(AsyncHTTPTestCase):
+    @classmethod
+    def setUpClass(cls):
+        prefix = "__TestEvaluationPlaneHandlerMaxRequestSize_"
+
+        # create config file
+        cls.config_file = tempfile.NamedTemporaryFile(
+            mode="w+t", prefix=prefix, suffix=".conf", delete=False
+        )
+        cls.config_file.write(
+            "[TabPy]\n"
+            "TABPY_MAX_REQUEST_SIZE_MB = 1"
+        )
+        cls.config_file.close()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.remove(cls.config_file.name)
+
+    def get_app(self):
+        self.app = TabPyApp(self.config_file.name)
+        return self.app._create_tornado_web_app()
+
+    def create_large_payload(self):
+        num_chars = 2 * 1024 * 1024 # 2MB Size
+        large_string = string.printable * (num_chars // len(string.printable))
+        large_string += string.printable[:num_chars % len(string.printable)]
+        payload = {
+            "data": { "_arg1": [1, large_string] },
+            "script": "return _arg1"
+        }
+        return json.dumps(payload).encode('utf-8')
+
+    def test_evaluation_payload_exceeds_max_request_size(self):
+        response = self.fetch(
+            "/evaluate",
+            method="POST",
+            body=self.create_large_payload()
+        )
+        self.assertEqual(413, response.code)
+
+    def test_evaluation_max_request_size_not_applied(self):
+        self.app.max_request_size = None
+        response = self.fetch(
+            "/evaluate",
+            method="POST",
+            body=self.create_large_payload()
+        )
+        self.assertEqual(200, response.code)
+        self.assertEqual(1, json.loads(response.body)[0])
+
+    def test_no_content_length_header_present(self):
+        response = self.fetch(
+            "/evaluate",
+            method="POST",
+            allow_nonstandard_methods=True
+        )
+        message = json.loads(response.body)["message"]
+        # Ensure it reaches script processing stage in EvaluationPlaneHandler.post
+        self.assertEqual("Error processing script", message)
+
 
 class TestEvaluationPlaneHandlerDefault(AsyncHTTPTestCase):
     @classmethod