@pytest.mark.parametrize(
    "sampler_cls,sampler_kwargs",
    [
        (optuna.samplers.RandomSampler, {}),
        (optuna.samplers.TPESampler, {"n_startup_trials": 1}),
        (
            optuna.samplers.GridSampler,
            {"search_space": {"value": numpy.linspace(0.0, 1.0, 10, endpoint=False).tolist()}},
        ),
        (optuna.samplers.CmaEsSampler, {"n_startup_trials": 1}),
    ],
)
def test_hyperband_filter_study(sampler_cls, sampler_kwargs):
    # type: () -> None
    def objective(trial: optuna.trial.Trial) -> float:
        return trial.suggest_uniform("value", 0.0, 1.0)
    n_trials = 10
    n_brackets = 2
    expected_n_trials_per_bracket = n_trials // n_brackets
    with mock.patch(
        "optuna.pruners.HyperbandPruner._get_bracket_id",
        new=mock.Mock(side_effect=lambda study, trial: trial.number % n_brackets),
    ):
        for method_name in [
            "infer_relative_search_space",
            "sample_relative",
            "sample_independent",
        ]:
            sampler = sampler_cls(**sampler_kwargs)
            pruner = optuna.pruners.HyperbandPruner(
                min_resource=MIN_RESOURCE,
                max_resource=MAX_RESOURCE,
                reduction_factor=REDUCTION_FACTOR,
            )
            with mock.patch(
                "optuna.samplers.{}.{}".format(sampler_cls.__name__, method_name),
                wraps=getattr(sampler, method_name),
            ) as method_mock:
                study = optuna.study.create_study(sampler=sampler, pruner=pruner)
                study.optimize(objective, n_trials=n_trials)
                args = method_mock.call_args[0]
                study = args[0]
                trials = study.get_trials()
                assert len(trials) == expected_n_trials_per_bracket