From d04101fb769c171fb435789e1d0edaa96dc10063 Mon Sep 17 00:00:00 2001
From: nabenabe0928 <shuhei.watanabe.utokyo@gmail.com>
Date: Tue, 26 Dec 2023 09:12:16 +0100
Subject: [PATCH] Add some tests

---
 python_tests/test_api.py | 73 +++++++++++++++-------------------------
 1 file changed, 28 insertions(+), 45 deletions(-)

diff --git a/python_tests/test_api.py b/python_tests/test_api.py
index e7210710..ebe5a0ef 100644
--- a/python_tests/test_api.py
+++ b/python_tests/test_api.py
@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import json
+import pytest
 import sys
 from unittest import TestCase
 
 import optuna
 from optuna import get_all_study_summaries
 from optuna.study import StudyDirection
+from optuna.trial import FrozenTrial
 from optuna_dashboard._app import create_app
 from optuna_dashboard._app import create_new_study
 from optuna_dashboard._note import note_str_key_prefix
@@ -44,70 +46,51 @@ class APITestCase(TestCase):
         study_summaries = json.loads(body)["study_summaries"]
         self.assertEqual(len(study_summaries), 2)
 
-    def test_get_study_details_without_after_param(self) -> None:
+    def run_get_study_details(
+        self, queries: dict[str, str] | None = None, expected_status: int = 200,
+    ) -> tuple[int, list[FrozenTrial]]:
         study = optuna.create_study()
         study_id = study._study_id
-        study.optimize(objective, n_trials=2)
+        study.optimize(objective, n_trials=10)
         app = create_app(study._storage)
 
         status, _, body = send_request(
             app,
             f"/api/studies/{study_id}",
             "GET",
+            queries=queries,
             content_type="application/json",
         )
-        self.assertEqual(status, 200)
-        all_trials = json.loads(body)["trials"]
-        self.assertEqual(len(all_trials), 2)
+        self.assertEqual(status, expected_status)
+        if expected_status == 400:
+            return []
+        else:
+            return json.loads(body)["trials"]
+
+    def test_get_study_details_without_after_param(self) -> None:
+        all_trials = self.run_get_study_details()
+        self.assertEqual(len(all_trials), 10)
 
     def test_get_study_details_with_after_param_partial(self) -> None:
-        study = optuna.create_study()
-        study_id = study._study_id
-        study.optimize(objective, n_trials=2)
-        app = create_app(study._storage)
+        all_trials = self.run_get_study_details({"after": "5"})
+        self.assertEqual(len(all_trials), 5)
 
-        status, _, body = send_request(
-            app,
-            f"/api/studies/{study_id}",
-            "GET",
-            queries={"after": "1"},
-            content_type="application/json",
-        )
-        self.assertEqual(status, 200)
-        all_trials = json.loads(body)["trials"]
-        self.assertEqual(len(all_trials), 1)
+    def test_get_study_details_with_params(self) -> None:
+        for after in [0, 5, 9, 10]:
+            for limit in [1, 2, 5, 10]:
+                trials = self.run_get_study_details({"after": str(after), "limit": str(limit)})
+                ans = list(range(after, min(10, after + limit)))
+                self.assertEqual([t["number"] for t in trials], ans)
 
     def test_get_study_details_with_after_param_full(self) -> None:
-        study = optuna.create_study()
-        study_id = study._study_id
-        study.optimize(objective, n_trials=2)
-        app = create_app(study._storage)
-
-        status, _, body = send_request(
-            app,
-            f"/api/studies/{study_id}",
-            "GET",
-            queries={"after": "2"},
-            content_type="application/json",
-        )
-        self.assertEqual(status, 200)
-        all_trials = json.loads(body)["trials"]
+        all_trials = self.run_get_study_details({"after": "10"})
         self.assertEqual(len(all_trials), 0)
 
     def test_get_study_details_with_after_param_illegal(self) -> None:
-        study = optuna.create_study()
-        study_id = study._study_id
-        study.optimize(objective, n_trials=2)
-        app = create_app(study._storage)
+        self.run_get_study_details({"after": "-1"}, expected_status=400)
 
-        status, _, body = send_request(
-            app,
-            f"/api/studies/{study_id}",
-            "GET",
-            queries={"after": "-1"},
-            content_type="application/json",
-        )
-        self.assertEqual(status, 400)
+    def test_get_study_details_with_limit_param_illegal(self) -> None:
+        self.run_get_study_details({"limit": "-1"}, expected_status=400)
 
     @pytest.mark.skipif(sys.version_info < (3, 8), reason="BoTorch dropped Python3.7 support")
     def test_get_best_trials_of_preferential_study(self) -> None: