""" 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