Function test_trials_command

tests/test_cli.py:533–606

(objective: Callable[[Trial], float], output_format: str | None)

@pytest.mark.parametrize("objective", (objective_func, objective_func_branched_search_space))
@output_formats
def test_trials_command(objective: Callable[[Trial], float], output_format: str | None) -> None:
    with NamedTemporaryFilePool() as fp, StorageSupplier("journal", file=fp) as storage:
        assert isinstance(storage, JournalStorage)
        storage_url = fp.name
        study_name = "test_study"
        n_trials = 10

        study = optuna.create_study(storage=storage, study_name=study_name)
        study.optimize(objective, n_trials=n_trials)
        attrs = (
            "number",
            "value",
            "datetime_start",
            "datetime_complete",
            "duration",
            "params",
            "user_attrs",
            "state",
        )

        # Run command.
        command = [
            "optuna",
            "trials",
            "--storage",
            storage_url,
            "--study-name",
            study_name,
        ]

        if output_format is not None:
            command += ["--format", output_format]

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

        assert len(trials) == n_trials

        df = study.trials_dataframe(attrs, multi_index=True)

        for i, trial in enumerate(trials):
            for key in df.columns:
                expected_value = df.loc[i][key]

                # The param may be NaN when the objective function has branched search space.
                if (
                    key[0] == "params"
                    and isinstance(expected_value, float)
                    and np.isnan(expected_value)
                ):
                    if output_format in (None, "table", "value"):
                        assert key[1] not in eval(trial["params"])
                    else:
                        assert key[1] not in trial["params"]
                    continue

                if key[1] == "":
                    value = trial[key[0]]
                else: