learningcircuit--local-deep-research/tests/benchmarks/optimization/test_optuna_optimizer_extra_coverage.py

"""
Extra coverage tests for benchmarks/optimization/optuna_optimizer.py.

Targets the 74 missing lines not covered by test_optuna_optimizer_coverage.py:
- _get_default_param_space structure
- _objective int/categorical param suggestion paths
- _objective with sanitize_data
- _run_experiment success with combined score
- _save_results with/without study, numpy arrays
- _create_visualizations paths (PLOTTING_AVAILABLE, trial counts)
- optimize() starting callback and no-callback paths
- _optimization_callback with study.best_value
"""

import numpy as np
from unittest.mock import Mock, patch

MODULE = "local_deep_research.benchmarks.optimization.optuna_optimizer"


def _make_optimizer(**kwargs):
    from local_deep_research.benchmarks.optimization.optuna_optimizer import (
        OptunaOptimizer,
    )

    defaults = {"base_query": "extra coverage query"}
    defaults.update(kwargs)
    return OptunaOptimizer(**defaults)


# ---------------------------------------------------------------------------
# _get_default_param_space
# ---------------------------------------------------------------------------


class TestGetDefaultParamSpace:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_param_space_contains_required_keys(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        assert "iterations" in space
        assert "questions_per_iteration" in space
        assert "search_strategy" in space
        assert "max_results" in space

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_iterations_is_int_type(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        assert space["iterations"]["type"] == "int"
        assert space["iterations"]["low"] == 1
        assert space["iterations"]["high"] == 5

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_search_strategy_is_categorical_with_choices(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        assert space["search_strategy"]["type"] == "categorical"
        assert "source-based" in space["search_strategy"]["choices"]
        assert "focused-iteration" in space["search_strategy"]["choices"]


# ---------------------------------------------------------------------------
# _objective – int and categorical suggestion paths
# ---------------------------------------------------------------------------


class TestObjectiveParamSuggestionTypes:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_int_param_suggested_with_step(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.suggest_int.return_value = 3
        mock_trial.suggest_categorical.return_value = "rapid"
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.5}
            param_space = {
                "iterations": {"type": "int", "low": 1, "high": 5, "step": 1}
            }
            optimizer._objective(mock_trial, param_space=param_space)
        mock_trial.suggest_int.assert_called_once_with(
            "iterations", 1, 5, step=1
        )

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_int_param_suggested_without_step(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 1
        mock_trial.suggest_int.return_value = 2
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.4}
            param_space = {"count": {"type": "int", "low": 1, "high": 10}}
            optimizer._objective(mock_trial, param_space=param_space)
        mock_trial.suggest_int.assert_called_once_with("count", 1, 10, step=1)

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_categorical_param_suggested(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 2
        mock_trial.suggest_categorical.return_value = "standard"
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.6}
            param_space = {
                "strategy": {
                    "type": "categorical",
                    "choices": ["standard", "rapid"],
                }
            }
            optimizer._objective(mock_trial, param_space=param_space)
        mock_trial.suggest_categorical.assert_called_once_with(
            "strategy", ["standard", "rapid"]
        )

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_returns_score_from_run_experiment(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.suggest_int.return_value = 2
        mock_trial.suggest_categorical.return_value = "iterdrag"
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.73}
            param_space = optimizer._get_default_param_space()
            result = optimizer._objective(mock_trial, param_space=param_space)
        assert result == 0.73

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_appends_to_trials_history(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 5
        mock_trial.suggest_int.return_value = 1
        mock_trial.suggest_categorical.return_value = "source_based"
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.55}
            param_space = optimizer._get_default_param_space()
            optimizer._objective(mock_trial, param_space=param_space)
        assert len(optimizer.trials_history) == 1
        assert optimizer.trials_history[0]["score"] == 0.55


# ---------------------------------------------------------------------------
# _save_results – sanitize_data path
# ---------------------------------------------------------------------------


class TestSaveResultsSanitizeData:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(f"{MODULE}.sanitize_data", side_effect=lambda x: x)
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_sanitize_data_called_during_save(
        self,
        mock_write_json,
        mock_sanitize,
        mock_joblib,
        mock_evaluator,
        tmp_path,
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = None
        optimizer.trials_history = [
            {"trial_number": 0, "score": 0.5, "params": {}}
        ]
        optimizer._save_results()
        mock_sanitize.assert_called()


# ---------------------------------------------------------------------------
# _run_experiment – combined score calculation
# ---------------------------------------------------------------------------


class TestRunExperimentCombinedScore:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_combined_score_with_quality_and_speed_weights(
        self, mock_profiler_cls, mock_evaluator
    ):
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.9,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance
        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 60.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer(
            metric_weights={"quality": 0.7, "speed": 0.3}
        )
        result = optimizer._run_experiment(
            {"iterations": 2, "questions_per_iteration": 2}
        )
        assert result["success"] is True
        assert result["quality_score"] == 0.9
        assert 0.0 <= result["score"] <= 1.0

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_run_experiment_includes_timing_info(
        self, mock_profiler_cls, mock_evaluator
    ):
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.7,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance
        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 45.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer()
        result = optimizer._run_experiment({"iterations": 1})
        assert "total_duration" in result
        assert result["total_duration"] == 45.0


# ---------------------------------------------------------------------------
# _save_results – edge cases
# ---------------------------------------------------------------------------


class TestSaveResultsEdgeCases:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_with_empty_trials_history(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = None
        optimizer.trials_history = []
        optimizer._save_results()
        mock_write_json.assert_called_once()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_numpy_array_converted(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = None
        optimizer.trials_history = [
            {
                "trial_number": 0,
                "score": np.float32(0.65),
                "params": {"max_results": np.int32(50)},
            }
        ]
        optimizer._save_results()
        written_data = mock_write_json.call_args_list[0][0][1]
        assert isinstance(written_data[0]["score"], float)

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_with_study_writes_best_params_json(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        mock_study = Mock()
        mock_study.best_params = {"iterations": 3, "max_results": 60}
        mock_study.best_value = 0.88
        mock_study.trials = [Mock(), Mock()]
        optimizer.study = mock_study
        optimizer.trials_history = []
        optimizer._save_results()
        # 2 JSON writes: trials + best params
        assert mock_write_json.call_count == 2
        # Also dumps the study via joblib
        mock_joblib.dump.assert_called_once()


# ---------------------------------------------------------------------------
# _create_visualizations
# ---------------------------------------------------------------------------


class TestCreateVisualizations:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", False)
    def test_create_visualizations_returns_early_without_plotting(
        self, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        optimizer.study = Mock()
        optimizer.study.trials = [Mock(), Mock()]
        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            optimizer._create_visualizations()
            mock_plot.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_create_visualizations_returns_early_when_no_study(
        self, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        optimizer.study = None
        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            optimizer._create_visualizations()
            mock_plot.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_create_visualizations_returns_early_fewer_than_2_trials(
        self, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_study.trials = [Mock()]  # only 1 trial
        optimizer.study = mock_study
        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            optimizer._create_visualizations()
            mock_plot.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    @patch(f"{MODULE}.plot_optimization_history")
    @patch(f"{MODULE}.plot_param_importances")
    @patch(f"{MODULE}.plot_contour")
    @patch(f"{MODULE}.plot_slice")
    def test_create_visualizations_calls_all_plots_with_sufficient_trials(
        self,
        mock_slice,
        mock_contour,
        mock_importances,
        mock_history,
        mock_evaluator,
        tmp_path,
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        mock_study = Mock()
        mock_study.trials = [Mock() for _ in range(5)]
        optimizer.study = mock_study
        for mock_fn in [
            mock_history,
            mock_importances,
            mock_contour,
            mock_slice,
        ]:
            mock_fig = Mock()
            mock_fn.return_value = mock_fig
        optimizer._create_visualizations()
        mock_history.assert_called_once_with(mock_study)


# ---------------------------------------------------------------------------
# optimize() – starting callback and study creation
# ---------------------------------------------------------------------------


class TestOptimizeStartingCallback:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_starting_callback_fired_before_optimize(
        self, mock_optuna, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        callback = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 2}
        mock_study.best_value = 0.5
        mock_study.trials = [Mock()]
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()
        optimizer = _make_optimizer(n_trials=1, progress_callback=callback)
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()
        starting_calls = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "starting"
        ]
        assert len(starting_calls) == 1
        assert starting_calls[0][0][2]["stage"] == "initialization"

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_no_callback_does_not_raise(
        self, mock_optuna, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.3
        mock_study.trials = []
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()
        optimizer = _make_optimizer(n_trials=1)
        assert optimizer.progress_callback is None
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            params, value = optimizer.optimize()
        assert params == {"iterations": 1}
        assert value == 0.3


# ---------------------------------------------------------------------------
# _optimization_callback – best_value logging path
# ---------------------------------------------------------------------------


class TestOptimizationCallbackBestValue:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_at_trial_1_does_not_save(self, mock_evaluator):
        """Trial 1 is not a multiple of 10, so no save is triggered."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_study.best_value = 0.77
        mock_trial = Mock()
        mock_trial.number = 1  # 1 % 10 != 0, no save
        with patch.object(optimizer, "_save_results") as mock_save:
            optimizer._optimization_callback(mock_study, mock_trial)
        mock_save.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_at_trial_10_triggers_save(self, mock_evaluator):
        """Trial 10 is a multiple of 10 and > 0, so save is triggered."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 10
        with (
            patch.object(optimizer, "_save_results") as mock_save,
            patch.object(optimizer, "_create_quick_visualizations") as mock_viz,
        ):
            optimizer._optimization_callback(mock_study, mock_trial)
        mock_save.assert_called_once()
        mock_viz.assert_called_once()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_at_multiple_of_10_triggers_save(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 30
        with (
            patch.object(optimizer, "_save_results") as mock_save,
            patch.object(optimizer, "_create_quick_visualizations") as mock_viz,
        ):
            optimizer._optimization_callback(mock_study, mock_trial)
        mock_save.assert_called_once()
        mock_viz.assert_called_once()


# ---------------------------------------------------------------------------
# metric_weights normalization
# ---------------------------------------------------------------------------


class TestMetricWeightsNormalization:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_weights_normalized_to_sum_one(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(
            metric_weights={"quality": 3.0, "speed": 1.0}
        )
        total = sum(optimizer.metric_weights.values())
        assert abs(total - 1.0) < 1e-9

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_benchmark_weights_stored(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        weights = {"simpleqa": 0.6, "browsecomp": 0.4}
        optimizer = _make_optimizer(benchmark_weights=weights)
        assert optimizer.benchmark_weights == weights