learningcircuit--local-deep-research/tests/benchmarks/test_optimization_api.py

"""
Tests for optimization API functions.

This module tests the convenience functions that wrap OptunaOptimizer
for different optimization strategies.
"""

from unittest.mock import MagicMock, patch


from local_deep_research.benchmarks.optimization.api import (
    optimize_for_efficiency,
    optimize_for_quality,
    optimize_for_speed,
    optimize_parameters,
)


class TestOptimizeParameters:
    """Tests for the optimize_parameters function."""

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_creates_optimizer_with_query(self, mock_optimizer_class):
        """Function creates optimizer with the provided query."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({"iterations": 3}, 0.85)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_parameters(query="test research query")

        mock_optimizer_class.assert_called_once()
        call_kwargs = mock_optimizer_class.call_args[1]
        assert call_kwargs["base_query"] == "test research query"

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_passes_all_parameters_to_optimizer(self, mock_optimizer_class):
        """Function passes all configuration parameters to optimizer."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_parameters(
            query="test query",
            output_dir="/custom/output",
            model_name="gpt-4",
            provider="openai",
            search_tool="google",
            temperature=0.5,
            n_trials=50,
            timeout=3600,
            n_jobs=4,
            study_name="custom_study",
            optimization_metrics=["quality"],
            metric_weights={"quality": 1.0},
            benchmark_weights={"simpleqa": 0.7, "browsecomp": 0.3},
        )

        call_kwargs = mock_optimizer_class.call_args[1]
        assert call_kwargs["output_dir"] == "/custom/output"
        assert call_kwargs["model_name"] == "gpt-4"
        assert call_kwargs["provider"] == "openai"
        assert call_kwargs["search_tool"] == "google"
        assert call_kwargs["temperature"] == 0.5
        assert call_kwargs["n_trials"] == 50
        assert call_kwargs["timeout"] == 3600
        assert call_kwargs["n_jobs"] == 4
        assert call_kwargs["study_name"] == "custom_study"
        assert call_kwargs["optimization_metrics"] == ["quality"]
        assert call_kwargs["metric_weights"] == {"quality": 1.0}
        assert call_kwargs["benchmark_weights"] == {
            "simpleqa": 0.7,
            "browsecomp": 0.3,
        }

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_calls_optimizer_optimize_method(self, mock_optimizer_class):
        """Function calls the optimizer's optimize method."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        param_space = {"iterations": {"type": "int", "low": 1, "high": 5}}
        optimize_parameters(query="test", param_space=param_space)

        mock_optimizer.optimize.assert_called_once_with(param_space)

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_returns_optimizer_result(self, mock_optimizer_class):
        """Function returns the result from optimizer."""
        expected_params = {"iterations": 3, "search_strategy": "rapid"}
        expected_score = 0.92
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = (expected_params, expected_score)
        mock_optimizer_class.return_value = mock_optimizer

        result_params, result_score = optimize_parameters(query="test")

        assert result_params == expected_params
        assert result_score == expected_score


class TestOptimizeForSpeed:
    """Tests for the optimize_for_speed function."""

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_uses_speed_focused_param_space(self, mock_optimizer_class):
        """Function uses a parameter space optimized for speed."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_speed(query="test")

        # Check the param_space passed to optimize()
        param_space = mock_optimizer.optimize.call_args[0][0]
        # Speed-focused should have limited iterations (max 3)
        assert param_space["iterations"]["high"] == 3
        assert param_space["questions_per_iteration"]["high"] == 3

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_uses_speed_focused_weights(self, mock_optimizer_class):
        """Function uses metric weights that prioritize speed."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_speed(query="test")

        call_kwargs = mock_optimizer_class.call_args[1]
        metric_weights = call_kwargs["metric_weights"]
        # Speed should be heavily weighted
        assert metric_weights["speed"] == 0.8
        assert metric_weights["quality"] == 0.2
        assert metric_weights["resource"] == 0.0

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_uses_speed_focused_search_strategies(self, mock_optimizer_class):
        """Function uses fast search strategies."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_speed(query="test")

        param_space = mock_optimizer.optimize.call_args[0][0]
        strategies = param_space["search_strategy"]["choices"]
        # Should include fast strategies
        assert "source-based" in strategies
        assert "focused-iteration" in strategies


class TestOptimizeForQuality:
    """Tests for the optimize_for_quality function."""

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_uses_quality_focused_weights(self, mock_optimizer_class):
        """Function uses metric weights that prioritize quality."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_quality(query="test")

        call_kwargs = mock_optimizer_class.call_args[1]
        metric_weights = call_kwargs["metric_weights"]
        # Quality should be heavily weighted
        assert metric_weights["quality"] == 0.9
        assert metric_weights["speed"] == 0.1
        assert metric_weights["resource"] == 0.0

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_uses_default_param_space(self, mock_optimizer_class):
        """Function passes None for param_space (uses default)."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_quality(query="test")

        # param_space should be None (will use optimizer's default)
        param_space = mock_optimizer.optimize.call_args[0][0]
        assert param_space is None

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_includes_quality_in_optimization_metrics(
        self, mock_optimizer_class
    ):
        """Function includes quality in optimization metrics."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_quality(query="test")

        call_kwargs = mock_optimizer_class.call_args[1]
        assert "quality" in call_kwargs["optimization_metrics"]


class TestOptimizeForEfficiency:
    """Tests for the optimize_for_efficiency function."""

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_uses_balanced_weights(self, mock_optimizer_class):
        """Function uses balanced metric weights."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_efficiency(query="test")

        call_kwargs = mock_optimizer_class.call_args[1]
        metric_weights = call_kwargs["metric_weights"]
        # Should balance all three metrics
        assert metric_weights["quality"] == 0.4
        assert metric_weights["speed"] == 0.3
        assert metric_weights["resource"] == 0.3

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_includes_resource_metric(self, mock_optimizer_class):
        """Function includes resource in optimization metrics."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_efficiency(query="test")

        call_kwargs = mock_optimizer_class.call_args[1]
        assert "resource" in call_kwargs["optimization_metrics"]

    @patch("local_deep_research.benchmarks.optimization.api.OptunaOptimizer")
    def test_includes_all_three_metrics(self, mock_optimizer_class):
        """Function optimizes for quality, speed, and resource."""
        mock_optimizer = MagicMock()
        mock_optimizer.optimize.return_value = ({}, 0.5)
        mock_optimizer_class.return_value = mock_optimizer

        optimize_for_efficiency(query="test")

        call_kwargs = mock_optimizer_class.call_args[1]
        metrics = call_kwargs["optimization_metrics"]
        assert "quality" in metrics
        assert "speed" in metrics
        assert "resource" in metrics