mirror of
https://github.com/wassname/optuna-dashboard.git
synced 2026-08-22 12:20:17 +08:00
83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
from optuna.storages import BaseStorage
|
|
from optuna_dashboard.artifact import _backend
|
|
import pytest
|
|
|
|
|
|
def test_get_artifact_path() -> None:
|
|
study = MagicMock(_study_id=0)
|
|
trial = MagicMock(_trial_id=0, study=study)
|
|
assert _backend.get_artifact_path(trial=trial, artifact_id="id0") == "/artifacts/0/0/id0"
|
|
|
|
|
|
def test_artifact_prefix() -> None:
|
|
actual = _backend._artifact_prefix(trial_id=0)
|
|
assert actual == "dashboard:artifacts:0:"
|
|
|
|
|
|
@pytest.fixture()
|
|
def init_storage_with_artifact_meta() -> BaseStorage:
|
|
from optuna import create_study
|
|
from optuna.storages import InMemoryStorage
|
|
|
|
storage = InMemoryStorage()
|
|
study = create_study(storage=storage)
|
|
|
|
study_system_attrs = {
|
|
"dashboard:artifacts:0:id0": '{"artifact_id": "id0", "filename": "foo.txt"}',
|
|
"dashboard:artifacts:0:id1": '{"artifact_id": "id1", "filename": "bar.txt"}',
|
|
"baz": "baz",
|
|
}
|
|
for key, value in study_system_attrs.items():
|
|
study.set_system_attr(key, value)
|
|
|
|
trial_system_attrs = {
|
|
"artifacts:id2": '{"artifact_id": "id2", "filename": "baz.txt"}',
|
|
"artifacts:id3": '{"artifact_id": "id3", "filename": "qux.txt"}',
|
|
}
|
|
for key, value in trial_system_attrs.items():
|
|
trial = study.ask()
|
|
trial.set_system_attr(key, value)
|
|
study.tell(trial, 0.0)
|
|
|
|
return storage
|
|
|
|
|
|
def test_get_artifact_meta(init_storage_with_artifact_meta: MagicMock) -> None:
|
|
storage = init_storage_with_artifact_meta
|
|
|
|
actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=0, artifact_id="id0")
|
|
assert actual == {"artifact_id": "id0", "filename": "foo.txt"}
|
|
|
|
actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=1, artifact_id="id3")
|
|
assert actual == {"artifact_id": "id3", "filename": "qux.txt"}
|
|
|
|
actual = _backend.get_artifact_meta(storage, study_id=0, trial_id=0, artifact_id="id4")
|
|
assert actual is None
|
|
|
|
|
|
def test_delete_all_artifacts(init_storage_with_artifact_meta: MagicMock) -> None:
|
|
backend = MagicMock()
|
|
storage = init_storage_with_artifact_meta
|
|
_backend.delete_all_artifacts(backend, storage, study_id=0)
|
|
|
|
assert backend.remove.call_args_list == [
|
|
(("id0",),),
|
|
(("id1",),),
|
|
(("id2",),),
|
|
(("id3",),),
|
|
]
|
|
|
|
|
|
def test_list_trial_artifacts(init_storage_with_artifact_meta: MagicMock) -> None:
|
|
storage = init_storage_with_artifact_meta
|
|
trial = MagicMock(_trial_id=0, system_attrs=storage.get_trial_system_attrs(0))
|
|
|
|
actual = _backend.list_trial_artifacts(storage.get_study_system_attrs(0), trial)
|
|
assert actual == [
|
|
{"artifact_id": "id0", "filename": "foo.txt"},
|
|
{"artifact_id": "id1", "filename": "bar.txt"},
|
|
{"artifact_id": "id2", "filename": "baz.txt"},
|
|
]
|