Update test_csv_download.py

This commit is contained in:
Hiroki Takizawa
2023-12-07 10:58:16 +09:00
committed by GitHub
parent 678b15e60c
commit 18d68ea919
+21 -48
View File
@@ -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)