From 18d68ea91959563d0364dc561ca4ee3576a1625a Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Thu, 7 Dec 2023 10:58:16 +0900
Subject: [PATCH] Update test_csv_download.py

---
 python_tests/test_csv_download.py | 69 ++++++++++---------------------
 1 file changed, 21 insertions(+), 48 deletions(-)

diff --git a/python_tests/test_csv_download.py b/python_tests/test_csv_download.py
index 7740ea83..e1ab12c9 100644
--- a/python_tests/test_csv_download.py
+++ b/python_tests/test_csv_download.py
@@ -8,6 +8,21 @@ import pytest
 from .wsgi_client import send_request
 
 
+def _validate_output(
+    storage: optuna.storages.BaseStorage,
+    correct_status: int,
+    study_id: int,
+) -> None:
+    app = create_app(storage)
+    status, _, _ = send_request(
+        app,
+        f"/csv/{study_id}",
+        "GET",
+        content_type="application/json",
+    )
+    assert status == correct_status
+
+
 def test_download_csv_no_trial() -> None:
     def objective(trial: optuna.Trial) -> float:
         x = trial.suggest_float("x", -100, 100)
@@ -17,42 +32,21 @@ def test_download_csv_no_trial() -> None:
     storage = optuna.storages.InMemoryStorage()
     study = optuna.create_study(storage=storage)
     study.optimize(objective, n_trials=0)
-    app = create_app(storage)
-    status, _, _ = send_request(
-        app,
-        "/csv/0",
-        "GET",
-        content_type="application/json",
-    )
-    assert status == 200
+    _validate_output(storage, 200, 0)
 
 
 def test_download_csv_all_waiting() -> None:
     storage = optuna.storages.InMemoryStorage()
     study = optuna.create_study(storage=storage)
     study.add_trial(optuna.trial.create_trial(state=TrialState.WAITING))
-    app = create_app(storage)
-    status, _, body = send_request(
-        app,
-        "/csv/0",
-        "GET",
-        content_type="application/json",
-    )
-    assert status == 200
+    _validate_output(storage, 200, 0)
 
 
 def test_download_csv_all_running() -> None:
     storage = optuna.storages.InMemoryStorage()
     study = optuna.create_study(storage=storage)
     study.add_trial(optuna.trial.create_trial(state=TrialState.RUNNING))
-    app = create_app(storage)
-    status, _, body = send_request(
-        app,
-        "/csv/0",
-        "GET",
-        content_type="application/json",
-    )
-    assert status == 200
+    _validate_output(storage, 200, 0)
 
 
 @pytest.mark.parametrize("study_id", [0, 1])
@@ -66,14 +60,7 @@ def test_download_csv_fail(study_id: int) -> None:
     study = optuna.create_study(storage=storage)
     optuna.logging.set_verbosity(optuna.logging.ERROR)
     study.optimize(objective, n_trials=10)
-    app = create_app(storage)
-    status, _, body = send_request(
-        app,
-        f"/csv/{id}",
-        "GET",
-        content_type="application/json",
-    )
-    assert status == (404 if id != 0 else 200)
+    _validate_output(storage, 404 if study_id != 0 else 200, study_id)
 
 
 @pytest.mark.parametrize("is_multi_obj", [True, False])
@@ -90,14 +77,7 @@ def test_download_csv_multi_obj(is_multi_obj: bool) -> None:
     study = optuna.create_study(storage=storage, directions=directions)
     optuna.logging.set_verbosity(optuna.logging.ERROR)
     study.optimize(objective, n_trials=10)
-    app = create_app(storage)
-    status, _, body = send_request(
-        app,
-        "/csv/0",
-        "GET",
-        content_type="application/json",
-    )
-    assert status == 200
+    _validate_output(storage, 200, 0)
 
 
 def test_download_csv_user_attr() -> None:
@@ -111,11 +91,4 @@ def test_download_csv_user_attr() -> None:
     study = optuna.create_study(storage=storage)
     optuna.logging.set_verbosity(optuna.logging.ERROR)
     study.optimize(objective, n_trials=10)
-    app = create_app(storage)
-    status, _, body = send_request(
-        app,
-        "/csv/0",
-        "GET",
-        content_type="application/json",
-    )
-    assert status == 200
+    _validate_output(storage, 200, 0)