From d433914a54e4b9290dd6669b23ee5557a8c1be17 Mon Sep 17 00:00:00 2001 From: c-bata Date: Wed, 4 Jan 2023 16:59:52 +0900 Subject: [PATCH] Add tests for save_note apis --- python_tests/test_note.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/python_tests/test_note.py b/python_tests/test_note.py index c53f1d0c..5e3c39bb 100644 --- a/python_tests/test_note.py +++ b/python_tests/test_note.py @@ -3,7 +3,10 @@ from __future__ import annotations from unittest import TestCase from unittest.mock import patch +import optuna from optuna_dashboard import _note as note +from optuna_dashboard import save_study_note +from optuna_dashboard import save_trial_note class NoteTestCase(TestCase): @@ -19,3 +22,26 @@ class NoteTestCase(TestCase): assert len(attrs) == attr_len actual = note.concat_body(attrs, None) assert dummy_body_str == actual + + def test_save_and_get_study_note(self) -> None: + study = optuna.create_study() + + for body, expected_ver in [("version 1", 1), ("version 2", 2)]: + with self.subTest(body): + save_study_note(study, body) + system_attrs = study._storage.get_study_system_attrs(study._study_id) + note_dict = note.get_note_from_system_attrs(system_attrs, None) + assert note_dict["body"] == body + assert note_dict["version"] == expected_ver + + def test_save_and_get_trial_note(self) -> None: + study = optuna.create_study() + trial = study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) + + for body, expected_ver in [("version 1", 1), ("version 2", 2)]: + with self.subTest(body): + save_trial_note(trial, body) + system_attrs = study._storage.get_study_system_attrs(study._study_id) + note_dict = note.get_note_from_system_attrs(system_attrs, trial._trial_id) + assert note_dict["body"] == body + assert note_dict["version"] == expected_ver