From c73abacef3c96569638b120d28a7965f220c79c6 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 6 Dec 2023 15:35:35 +0900
Subject: [PATCH] Add files via upload

---
 python_tests/test_csv_download.py | 73 +++++++++++++++++++++++++++++++
 1 file changed, 73 insertions(+)
 create mode 100644 python_tests/test_csv_download.py

diff --git a/python_tests/test_csv_download.py b/python_tests/test_csv_download.py
new file mode 100644
index 00000000..a820f37b
--- /dev/null
+++ b/python_tests/test_csv_download.py
@@ -0,0 +1,73 @@
+import optuna
+from optuna_dashboard._app import create_app
+import pytest
+
+from .wsgi_client import send_request
+
+
+@pytest.mark.parametrize("id", [0, 1])
+def test_download_csv_fail(id: int) -> None:
+    def objective(trial: optuna.Trial) -> float:
+        x = trial.suggest_float("x", -100, 100)
+        y = trial.suggest_categorical("y", [-1, 0, 1])
+        return x**2 + y
+
+    storage = optuna.storages.InMemoryStorage()
+    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)
+
+
+@pytest.mark.parametrize("is_multi_obj", [True, False])
+def test_download_csv_multi_obj(is_multi_obj: bool) -> None:
+    def objective(trial: optuna.Trial) -> float:
+        x = trial.suggest_float("x", -100, 100)
+        y = trial.suggest_categorical("y", [-1, 0, 1])
+        if is_multi_obj:
+            return x**2, y
+        return x**2 + y
+
+    storage = optuna.storages.InMemoryStorage()
+    if is_multi_obj:
+        study = optuna.create_study(storage=storage, directions=["minimize", "minimize"])
+    else:
+        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/{0}",
+        "GET",
+        content_type="application/json",
+    )
+    assert status == 200
+
+
+def test_download_csv_user_attr() -> None:
+    def objective(trial: optuna.Trial) -> float:
+        x = trial.suggest_float("x", -100, 100)
+        y = trial.suggest_categorical("y", [-1, 0, 1])
+        trial.set_user_attr("abs_y", abs(y))
+        return x**2 + y
+
+    storage = optuna.storages.InMemoryStorage()
+    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/{0}",
+        "GET",
+        content_type="application/json",
+    )
+    assert status == 200