Function test_studies_command

tests/test_cli.py:375–438

(output_format: str | None)

@output_formats
def test_studies_command(output_format: str | None) -> None:
    with NamedTemporaryFilePool() as fp, StorageSupplier("journal", file=fp) as storage:
        assert isinstance(storage, JournalStorage)
        storage_url = fp.name

        # First study.
        study_1 = optuna.create_study(storage=storage)

        # Run command.
        command = ["optuna", "studies", "--storage", storage_url]
        if output_format is not None:
            command += ["--format", output_format]

        studies = _get_output(command, output_format or "table")

        expected_keys = ["name", "direction", "n_trials", "datetime_start"]

        # Check user_attrs are not printed.
        if output_format in (None, "table", "value"):
            assert list(studies[0].keys()) == expected_keys
        else:
            assert set(studies[0].keys()) == set(expected_keys)

        # Add a second study.
        study_2 = optuna.create_study(
            storage=storage, study_name="study_2", directions=["minimize", "maximize"]
        )
        study_2.optimize(objective_func_multi_objective, n_trials=10)
        study_2.set_user_attr("key_1", "value_1")
        study_2.set_user_attr("key_2", "value_2")

        # Run command again to include second study.
        studies = _get_output(command, output_format or "table")

        expected_keys = ["name", "direction", "n_trials", "datetime_start", "user_attrs"]

        assert len(studies) == 2
        for study in studies:
            if output_format in (None, "table", "value"):
                assert list(study.keys()) == expected_keys
            else:
                assert set(study.keys()) == set(expected_keys)

        # Check study_name, direction, n_trials and user_attrs for the first study.
        assert studies[0]["name"] == study_1.study_name
        if output_format in (None, "table", "value"):
            assert studies[0]["n_trials"] == "0"
            assert eval(studies[0]["direction"]) == ("MINIMIZE",)
            assert eval(studies[0]["user_attrs"]) == {}
        else:
            assert studies[0]["n_trials"] == 0
            assert studies[0]["direction"] == ["MINIMIZE"]
            assert studies[0]["user_attrs"] == {}

        # Check study_name, direction, n_trials and user_attrs for the second study.
        assert studies[1]["name"] == study_2.study_name
        if output_format in (None, "table", "value"):
            assert studies[1]["n_trials"] == "10"