Merge pull request #6 from c-bata/refactor-wsgi-utils

Refactor wsgi utils

tests/test_api.py

@@ -1,23 +1,9 @@
 import json
-from typing import Dict, Optional, Any
 from unittest import TestCase
 
 import optuna
-from .wsgi_utils import create_wsgi_env, send_request, WSGIEnv
 from optuna_dashboard.app import create_app
-
-
-def create_json_api_wsgi_env(
-    path: str,
-    method: str,
-    json_body: Optional[Dict[str, Any]] = None,
-    content_type: str = "application/json",
-    headers: Optional[Dict[str, str]] = None,
-) -> WSGIEnv:
-    body = json.dumps(json_body) if json_body else ""
-    return create_wsgi_env(
-        path, method, body, content_type=content_type, headers=headers
-    )
+from .wsgi_client import send_request
 
 
 class APITestCase(TestCase):
@@ -27,12 +13,13 @@ class APITestCase(TestCase):
         storage.create_new_study("foo2")
 
         app = create_app(storage)
-        env = create_json_api_wsgi_env(
+        status, _, body = send_request(
+            app,
             "/api/studies/",
             "GET",
+            content_type="application/json",
         )
-        status, _, body = send_request(app, env)
-        self.assertEqual(status, "200 OK")
+        self.assertEqual(status, 200)
         study_summaries = json.loads(body)["study_summaries"]
         self.assertEqual(len(study_summaries), 2)
 
@@ -41,16 +28,18 @@ class APITestCase(TestCase):
         self.assertEqual(len(storage.get_all_study_summaries()), 0)
 
         app = create_app(storage)
-        env = create_json_api_wsgi_env(
+        request_body = {
+            "study_name": "foo",
+            "direction": "minimize",
+        }
+        status, _, _ = send_request(
+            app,
             "/api/studies",
             "POST",
-            json_body={
-                "study_name": "foo",
-                "direction": "minimize",
-            },
+            content_type="application/json",
+            body=json.dumps(request_body),
         )
-        status, _, _ = send_request(app, env)
-        self.assertEqual(status, "201 Created")
+        self.assertEqual(status, 201)
         self.assertEqual(len(storage.get_all_study_summaries()), 1)
 
     def test_create_study_duplicated(self) -> None:
@@ -59,16 +48,18 @@ class APITestCase(TestCase):
         self.assertEqual(len(storage.get_all_study_summaries()), 1)
 
         app = create_app(storage)
-        env = create_json_api_wsgi_env(
+        request_body = {
+            "study_name": "foo",
+            "direction": "minimize",
+        }
+        status, _, _ = send_request(
+            app,
             "/api/studies",
             "POST",
-            json_body={
-                "study_name": "foo",
-                "direction": "minimize",
-            },
+            content_type="application/json",
+            body=json.dumps(request_body),
         )
-        status, _, _ = send_request(app, env)
-        self.assertEqual(status, "400 Bad Request")
+        self.assertEqual(status, 400)
         self.assertEqual(len(storage.get_all_study_summaries()), 1)
 
     def test_delete_study(self) -> None:
@@ -78,20 +69,39 @@ class APITestCase(TestCase):
         self.assertEqual(len(storage.get_all_study_summaries()), 2)
 
         app = create_app(storage)
-        env = create_json_api_wsgi_env(
+        status, _, _ = send_request(
+            app,
             "/api/studies/1",
             "DELETE",
+            content_type="application/json",
         )
-        status, _, _ = send_request(app, env)
-        self.assertEqual(status, "204 No Content")
+        self.assertEqual(status, 204)
         self.assertEqual(len(storage.get_all_study_summaries()), 1)
 
     def test_delete_study_not_found(self) -> None:
         storage = optuna.storages.InMemoryStorage()
         app = create_app(storage)
-        env = create_json_api_wsgi_env(
+        status, _, _ = send_request(
+            app,
             "/api/studies/1",
             "DELETE",
+            content_type="application/json",
         )
-        status, _, _ = send_request(app, env)
-        self.assertEqual(status, "404 Not Found")
+        self.assertEqual(status, 404)
+
+
+class BottleRequestHookTestCase(TestCase):
+    def test_ignore_trailing_slashes(self) -> None:
+        storage = optuna.storages.InMemoryStorage()
+        app = create_app(storage)
+
+        endpoints = ["/api/studies", "/api/studies/"]
+        for endpoint in endpoints:
+            with self.subTest(msg=endpoint):
+                status, _, body = send_request(
+                    app,
+                    endpoint,
+                    "GET",
+                    content_type="application/json",
+                )
+                self.assertEqual(status, 200)

tests/wsgi_client.py

@@ -9,56 +9,61 @@ WSGIApp = Callable[[WSGIEnv, StartResponse], Iterable[bytes]]
 def create_wsgi_env(
     path: str,
     method: str,
-    body: Union[str, bytes] = b"",
-    queries: Optional[Dict[str, str]] = None,
-    headers: Optional[Dict[str, str]] = None,
-    content_type: str = "text/plain; charset=utf-8",
+    content_type: str,
+    body: bytes,
+    queries: Dict[str, str],
+    headers: Dict[str, str],
 ) -> WSGIEnv:
-    request_method = method.upper()
-    bytes_body = body if isinstance(body, bytes) else body.encode("utf-8")
-    wsgi_input = io.BytesIO(bytes_body)
-    content_length = len(body)
-
     # 'key1=value1&key2=value2'
-    query_string = "&".join([f"{k}={v}" for k, v in queries.items()]) if queries else ""
+    query_string = "&".join([f"{k}={v}" for k, v in queries.items()])
 
     # See https://www.python.org/dev/peps/pep-3333/#environ-variables
     env = {
         "PATH_INFO": path,
-        "REQUEST_METHOD": request_method,
+        "REQUEST_METHOD": method.upper(),
         "SCRIPT_NAME": "",
         "QUERY_STRING": query_string,
         "CONTENT_TYPE": content_type,
-        "CONTENT_LENGTH": content_length,
+        "CONTENT_LENGTH": len(body),
         "SERVER_PROTOCOL": "http",
         "SERVER_NAME": "localhost",
-        "wsgi.input": wsgi_input,
+        "wsgi.input": io.BytesIO(body),
         "wsgi.version": (1, 0),
         "wsgi.errors": io.StringIO(""),
         "wsgi.multithread": True,
         "wsgi.multitprocess": True,
         "wsgi.run_once": False,
     }
-    headers = headers or {}
     for k, v in headers.items():
         env[f"HTTP_{k.upper()}"] = v
     return env
 
 
 def send_request(
-    app: WSGIApp, env: WSGIEnv
-) -> Tuple[str, List[Tuple[str, str]], bytes]:
+    app: WSGIApp,
+    path: str,
+    method: str,
+    body: Union[str, bytes] = b"",
+    queries: Optional[Dict[str, str]] = None,
+    headers: Optional[Dict[str, str]] = None,
+    content_type: str = "text/plain; charset=utf-8",
+) -> Tuple[int, List[Tuple[str, str]], bytes]:
     status: str = ""
-    headers: List[Tuple[str, str]] = []
+    response_headers: List[Tuple[str, str]] = []
 
     def start_response(status_: str, headers_: List[Tuple[str, str]]) -> None:
-        nonlocal status, headers
+        nonlocal status, response_headers
         status = status_
-        headers = headers_
+        response_headers = headers_
 
+    bytes_body = body if isinstance(body, bytes) else body.encode("utf-8")
+    headers = headers or {}
+    queries = queries or {}
+    env = create_wsgi_env(path, method, content_type, bytes_body, queries, headers)
     body = b""
     iterable_body = app(env, start_response)
     for b in iterable_body:
         body += b
 
-    return status, headers, body
+    status_code = int(status.split()[0])
+    return status_code, response_headers, body