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

"""
Coverage tests for benchmarks/optimization/optuna_optimizer.py.
"""

import numpy as np
import pytest
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": "coverage test query"}
    defaults.update(kwargs)
    return OptunaOptimizer(**defaults)


class TestObjectiveFloatParamSuggestion:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_float_param_with_log_scale(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.suggest_float.return_value = 0.01
        param_space = {
            "lr": {"type": "float", "low": 0.0001, "high": 1.0, "log": True}
        }
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.77}
            score = optimizer._objective(mock_trial, param_space=param_space)
        mock_trial.suggest_float.assert_called_once_with(
            "lr", 0.0001, 1.0, step=None, log=True
        )
        assert score == 0.77

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_float_param_with_step_no_log(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_trial = Mock()
        mock_trial.number = 1
        mock_trial.suggest_float.return_value = 0.5
        param_space = {
            "dropout": {"type": "float", "low": 0.0, "high": 1.0, "step": 0.1}
        }
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.65}
            score = optimizer._objective(mock_trial, param_space=param_space)
        mock_trial.suggest_float.assert_called_once_with(
            "dropout", 0.0, 1.0, step=0.1, log=False
        )
        assert score == 0.65


class TestObjectiveProgressCallbacks:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_trial_started_then_completed(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        callback = Mock()
        optimizer = _make_optimizer(progress_callback=callback, n_trials=5)
        mock_trial = Mock()
        mock_trial.number = 2
        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.88}
            param_space = optimizer._get_default_param_space()
            optimizer._objective(mock_trial, param_space=param_space)
        stages = [c[0][2]["stage"] for c in callback.call_args_list]
        assert "trial_started" in stages
        assert "trial_completed" in stages
        completed_call = [
            c
            for c in callback.call_args_list
            if c[0][2]["stage"] == "trial_completed"
        ][0]
        assert completed_call[0][2]["score"] == 0.88

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_trial_error_on_exception(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        callback = Mock()
        optimizer = _make_optimizer(progress_callback=callback, n_trials=5)
        mock_trial = Mock()
        mock_trial.number = 3
        mock_trial.suggest_int.return_value = 1
        mock_trial.suggest_categorical.return_value = "standard"
        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.side_effect = RuntimeError("timeout")
            param_space = optimizer._get_default_param_space()
            score = optimizer._objective(mock_trial, param_space=param_space)
        assert score == float("-inf")
        error_calls = [
            c
            for c in callback.call_args_list
            if c[0][2].get("stage") == "trial_error"
        ]
        assert len(error_calls) == 1
        assert "timeout" in error_calls[0][0][2]["error"]


class TestOptimizeKeyboardInterrupt:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_keyboard_interrupt_saves_and_calls_callback(
        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.45
        mock_study.trials = [Mock(), Mock(), Mock()]
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study
        optimizer = _make_optimizer(n_trials=20, progress_callback=callback)
        with (
            patch.object(optimizer, "_save_results") as mock_save,
            patch.object(optimizer, "_create_visualizations") as mock_viz,
        ):
            best_params, best_value = optimizer.optimize()
        mock_save.assert_called_once()
        mock_viz.assert_called_once()
        assert best_params == {"iterations": 2}
        assert best_value == 0.45
        interrupted = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "interrupted"
        ]
        assert len(interrupted) == 1
        assert interrupted[0][0][2]["trials_completed"] == 3

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_keyboard_interrupt_without_callback(
        self, mock_optuna, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.1
        mock_study.trials = []
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study
        optimizer = _make_optimizer(n_trials=5)
        assert optimizer.progress_callback is None
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            best_params, best_value = optimizer.optimize()
        assert best_params == {"iterations": 1}


class TestOptimizeCompletionCallback:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_completion_callback_includes_best_params_and_value(
        self, mock_optuna, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        callback = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 4, "max_results": 80}
        mock_study.best_value = 0.93
        mock_study.trials = [Mock(), Mock()]
        mock_optuna.create_study.return_value = mock_study
        optimizer = _make_optimizer(n_trials=2, progress_callback=callback)
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()
        completed = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "completed"
            and c[0][2].get("stage") == "finished"
        ]
        assert len(completed) == 1
        info = completed[0][0][2]
        assert info["best_params"] == {"iterations": 4, "max_results": 80}
        assert info["best_value"] == 0.93
        assert info["trials_completed"] == 2


class TestOptimizationCallbackStoresTrial:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_saves_at_trial_20(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 20
        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_no_save_at_trial_5(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 5
        with patch.object(optimizer, "_save_results") as mock_save:
            optimizer._optimization_callback(mock_study, mock_trial)
        mock_save.assert_not_called()


class TestCreateQuickVisualizations:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    @patch(f"{MODULE}.plot_optimization_history")
    def test_quick_viz_with_sufficient_trials(
        self, mock_plot_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(), Mock(), Mock()]
        optimizer.study = mock_study
        mock_fig = Mock()
        mock_plot_history.return_value = mock_fig
        optimizer._create_quick_visualizations()
        mock_plot_history.assert_called_once_with(mock_study)
        mock_fig.write_image.assert_called_once()

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

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", False)
    def test_quick_viz_returns_early_without_matplotlib(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_quick_visualizations()
            mock_plot.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_quick_viz_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_quick_visualizations()
            mock_plot.assert_not_called()


class TestSaveResultsNumpyConversion:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_numpy_int64_top_level_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": np.int64(0),
                "score": np.float64(0.92),
                "params": {"iterations": np.int64(3)},
            }
        ]
        optimizer._save_results()
        assert mock_write_json.call_count == 1
        written_data = mock_write_json.call_args_list[0][0][1]
        assert isinstance(written_data[0]["trial_number"], float)
        assert isinstance(written_data[0]["score"], float)
        assert isinstance(written_data[0]["params"]["iterations"], float)

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_numpy_float64_in_result_dict_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,
                "result": {
                    "quality_score": np.float64(0.85),
                    "speed_score": np.float64(0.72),
                },
                "params": {},
                "score": 0.8,
            }
        ]
        optimizer._save_results()
        written_data = mock_write_json.call_args_list[0][0][1]
        result_dict = written_data[0]["result"]
        assert isinstance(result_dict["quality_score"], float)
        assert isinstance(result_dict["speed_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_saves_best_params_and_pkl(
        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}
        mock_study.best_value = 0.91
        mock_study.trials = [Mock()]
        optimizer.study = mock_study
        optimizer.trials_history = []
        optimizer._save_results()
        assert mock_write_json.call_count == 2
        mock_joblib.dump.assert_called_once()


class TestRunExperimentErrorPaths:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_evaluator_error_returns_failure_stops_profiler(
        self, mock_profiler_cls, mock_evaluator
    ):
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.side_effect = ValueError("bad config")
        mock_evaluator.return_value = mock_eval_instance
        mock_profiler = Mock()
        mock_profiler_cls.return_value = mock_profiler
        optimizer = _make_optimizer()
        result = optimizer._run_experiment(
            {"iterations": 2, "questions_per_iteration": 1}
        )
        assert result["success"] is False
        assert result["score"] == 0.0
        assert "bad config" in result["error"]
        mock_profiler.start.assert_called_once()
        mock_profiler.stop.assert_called_once()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_profiler_get_summary_error_caught(
        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.side_effect = RuntimeError("profiler broken")
        mock_profiler_cls.return_value = mock_profiler
        optimizer = _make_optimizer()
        result = optimizer._run_experiment({"iterations": 1})
        assert result["success"] is False
        assert result["score"] == 0.0
        assert "profiler broken" in result["error"]
        assert mock_profiler.stop.call_count >= 1

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_successful_experiment_returns_all_fields(
        self, mock_profiler_cls, mock_evaluator
    ):
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.85,
            "benchmark_results": {"simpleqa": {"accuracy": 0.85}},
        }
        mock_evaluator.return_value = mock_eval_instance
        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 120.0}
        mock_profiler_cls.return_value = mock_profiler
        optimizer = _make_optimizer(
            metric_weights={"quality": 0.6, "speed": 0.4}
        )
        result = optimizer._run_experiment(
            {
                "iterations": 2,
                "questions_per_iteration": 3,
                "search_strategy": "iterdrag",
                "max_results": 50,
            }
        )
        assert result["success"] is True
        assert result["quality_score"] == 0.85
        assert result["speed_score"] == pytest.approx(2 / 3, abs=0.01)
        assert result["total_duration"] == 120.0
        assert "score" in result
        assert "benchmark_results" in result