From b2dc9d6c9553953975138388a74d7ebf92ecea0d Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Sun, 24 Apr 2022 13:32:16 +0900
Subject: [PATCH] Fix tests with Optuna v3

---
 .github/workflows/python-tests.yml | 15 +++++++++++++++
 python_tests/test_api.py           | 27 ++++++++++++++++++++-------
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index 5ae71657..6a8c8c88 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -42,3 +42,18 @@ jobs:
           python -m pip install --progress-bar off --upgrade pip setuptools
           pip install --progress-bar off .
       - run: python -m unittest
+  test-with-optuna-master:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup Python 3.9
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.9'
+          architecture: x64
+      - name: Install dependencies
+        run: |
+          python -m pip install --progress-bar off --upgrade pip setuptools
+          pip install --progress-bar off .
+          python -m pip install --progress-bar off --upgrade git+https://github.com/optuna/optuna.git
+      - run: python -m unittest
diff --git a/python_tests/test_api.py b/python_tests/test_api.py
index b3c82c6d..2060dcee 100644
--- a/python_tests/test_api.py
+++ b/python_tests/test_api.py
@@ -1,13 +1,26 @@
 import json
+from typing import List
 from unittest import TestCase
 
 import optuna
+from optuna.storages import BaseStorage
+from optuna.study import StudySummary
+from optuna.version import __version__ as optuna_ver
+from packaging import version
 
 from optuna_dashboard._app import create_app
 
 from .wsgi_client import send_request
 
 
+def get_all_study_summaries(storage: BaseStorage) -> List[StudySummary]:
+    if version.parse(optuna_ver) >= version.Version("3.0.0b0.dev"):
+        summaries = storage.get_all_study_summaries(include_best_trial=True)  # type: ignore
+    else:
+        summaries = storage.get_all_study_summaries()  # type: ignore
+    return summaries
+
+
 def objective(trial: optuna.trial.Trial) -> float:
     x = trial.suggest_float("x", -1, 1)
     return x
@@ -103,7 +116,7 @@ class APITestCase(TestCase):
         ]:
             with self.subTest(name):
                 storage = optuna.storages.InMemoryStorage()
-                self.assertEqual(len(storage.get_all_study_summaries()), 0)
+                self.assertEqual(len(get_all_study_summaries(storage)), 0)
 
                 app = create_app(storage)
                 request_body = {
@@ -120,14 +133,14 @@ class APITestCase(TestCase):
                 self.assertEqual(status, expected_status)
 
                 if expected_status == 201:
-                    self.assertEqual(len(storage.get_all_study_summaries()), 1)
+                    self.assertEqual(len(get_all_study_summaries(storage)), 1)
                 else:
-                    self.assertEqual(len(storage.get_all_study_summaries()), 0)
+                    self.assertEqual(len(get_all_study_summaries(storage)), 0)
 
     def test_create_study_duplicated(self) -> None:
         storage = optuna.storages.InMemoryStorage()
         storage.create_new_study("foo")
-        self.assertEqual(len(storage.get_all_study_summaries()), 1)
+        self.assertEqual(len(get_all_study_summaries(storage)), 1)
 
         app = create_app(storage)
         request_body = {
@@ -142,13 +155,13 @@ class APITestCase(TestCase):
             body=json.dumps(request_body),
         )
         self.assertEqual(status, 400)
-        self.assertEqual(len(storage.get_all_study_summaries()), 1)
+        self.assertEqual(len(get_all_study_summaries(storage)), 1)
 
     def test_delete_study(self) -> None:
         storage = optuna.storages.InMemoryStorage()
         storage.create_new_study("foo1")
         storage.create_new_study("foo2")
-        self.assertEqual(len(storage.get_all_study_summaries()), 2)
+        self.assertEqual(len(get_all_study_summaries(storage)), 2)
 
         app = create_app(storage)
         status, _, _ = send_request(
@@ -158,7 +171,7 @@ class APITestCase(TestCase):
             content_type="application/json",
         )
         self.assertEqual(status, 204)
-        self.assertEqual(len(storage.get_all_study_summaries()), 1)
+        self.assertEqual(len(get_all_study_summaries(storage)), 1)
 
     def test_delete_study_not_found(self) -> None:
         storage = optuna.storages.InMemoryStorage()