comet-ml--opik/sdks/opik_optimizer/tests/unit/tools/test_wikipedia.py

"""Unit tests for Wikipedia search tools."""

from typing import Any

from pathlib import Path
import builtins
import sys
from types import SimpleNamespace
from unittest.mock import Mock, patch

import pytest

from opik_optimizer.utils.tools.wikipedia import (
    search_wikipedia,
    _search_wikipedia_api,
    _search_wikipedia_colbert,
    _search_wikipedia_bm25,
    _download_bm25_index,
)

pytestmark = pytest.mark.usefixtures("suppress_expected_optimizer_warnings")


class TestSearchWikipediaAPI:
    """Tests for Wikipedia API search."""

    @patch("opik_optimizer.utils.tools.wikipedia.requests.get")
    def test_api_search_success(self, mock_get: Mock) -> None:
        """Test successful API search."""
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "query": {
                "search": [
                    {
                        "title": "Python (programming language)",
                        "snippet": "Python is a <span>high-level</span> programming language",
                    }
                ]
            }
        }
        mock_get.return_value = mock_response

        results = _search_wikipedia_api("Python", k=1)

        assert len(results) == 1
        assert "Python (programming language)" in results[0]
        assert "high-level" in results[0]
        # HTML tags should be stripped
        assert "<span>" not in results[0]

    @patch("opik_optimizer.utils.tools.wikipedia.requests.get")
    def test_api_search_no_results(self, mock_get: Mock) -> None:
        """Test API search with no results."""
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"query": {"search": []}}
        mock_get.return_value = mock_response

        results = _search_wikipedia_api("nonexistentquery12345", k=1)

        assert len(results) == 1
        assert "No Wikipedia results found" in results[0]

    @patch("opik_optimizer.utils.tools.wikipedia.requests.get")
    def test_api_search_http_error(self, mock_get: Mock) -> None:
        """Test API search with HTTP error."""
        mock_response = Mock()
        mock_response.status_code = 500
        mock_get.return_value = mock_response

        with pytest.raises(Exception, match="Search API returned status 500"):
            _search_wikipedia_api("test", k=1)

    @patch("opik_optimizer.utils.tools.wikipedia.requests.get")
    def test_api_search_network_error(self, mock_get: Mock) -> None:
        """Test API search with network error."""
        mock_get.side_effect = Exception("Network error")

        with pytest.raises(Exception, match="Wikipedia API request failed"):
            _search_wikipedia_api("test", k=1)


class TestSearchWikipediaColBERT:
    """Tests for ColBERT search."""

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_colbert_not_available_fallback(self, mock_api: Mock) -> None:
        """Test fallback to API when ColBERT not available."""
        mock_api.return_value = ["fallback result"]

        # ColBERT module won't be available in most test environments
        results = _search_wikipedia_colbert("test", k=1)

        assert results == ["fallback result"]
        mock_api.assert_called_once_with("test", k=1)


class TestSearchWikipediaBM25:
    """Tests for BM25 search."""

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_bm25_not_available_fallback(self, mock_api: Mock) -> None:
        """Test fallback to API when BM25 not available."""
        mock_api.return_value = ["fallback result"]

        # bm25s module won't be available in most test environments
        results = _search_wikipedia_bm25("test", k=1, index_dir="/tmp/test")

        assert results == ["fallback result"]
        mock_api.assert_called_once_with("test", k=1)

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_bm25_no_index_fallback(self, mock_api: Mock) -> None:
        """Test fallback when no index provided."""
        mock_api.return_value = ["fallback result"]

        results = _search_wikipedia_bm25("test", k=1, index_dir=None, hf_repo=None)

        assert results == ["fallback result"]
        mock_api.assert_called_once_with("test", k=1)

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_bm25_index_not_found_fallback(self, mock_api: Mock) -> None:
        """Test fallback when index directory doesn't exist."""
        mock_api.return_value = ["fallback result"]

        # Test will naturally fall back because /nonexistent/path doesn't exist
        # The import will succeed, but the path check will fail
        results = _search_wikipedia_bm25("test", k=1, index_dir="/nonexistent/path")

        assert results == ["fallback result"]


class TestSearchWikipediaUnified:
    """Tests for unified search_wikipedia function."""

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_default_api_search(self, mock_api: Mock) -> None:
        """Test default search type is API."""
        mock_api.return_value = ["result"]

        results = search_wikipedia("test")

        assert results == ["result"]
        mock_api.assert_called_once_with("test", k=3)

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_explicit_api_search(self, mock_api: Mock) -> None:
        """Test explicit API search type."""
        mock_api.return_value = ["result"]

        results = search_wikipedia("test", search_type="api", k=5)

        assert results == ["result"]
        mock_api.assert_called_once_with("test", k=5)

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_colbert")
    def test_colbert_search(self, mock_colbert: Mock) -> None:
        """Test ColBERT search type."""
        mock_colbert.return_value = ["result"]

        results = search_wikipedia("test", search_type="colbert", k=5)

        assert results == ["result"]
        mock_colbert.assert_called_once_with("test", k=5)

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_bm25")
    def test_bm25_search(self, mock_bm25: Mock) -> None:
        """Test BM25 search type."""
        mock_bm25.return_value = ["result"]

        results = search_wikipedia(
            "test", search_type="bm25", k=5, bm25_index_dir="/tmp/index"
        )

        assert results == ["result"]
        mock_bm25.assert_called_once_with(
            "test", k=5, index_dir="/tmp/index", hf_repo=None
        )

    def test_invalid_search_type(self) -> None:
        """Test invalid search type raises error."""
        with pytest.raises(ValueError, match="Invalid search_type.*Must be"):
            search_wikipedia("test", search_type="invalid")  # type: ignore

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_api")
    def test_backward_compat_use_api_true(self, mock_api: Mock) -> None:
        """Test backward compatibility with use_api=True."""
        mock_api.return_value = ["result"]

        with pytest.warns(DeprecationWarning, match="use_api parameter is deprecated"):
            results = search_wikipedia("test", use_api=True)

        assert results == ["result"]
        mock_api.assert_called_once()

    @patch("opik_optimizer.utils.tools.wikipedia._search_wikipedia_colbert")
    def test_backward_compat_use_api_false(self, mock_colbert: Mock) -> None:
        """Test backward compatibility with use_api=False."""
        mock_colbert.return_value = ["result"]

        with pytest.warns(DeprecationWarning, match="use_api parameter is deprecated"):
            results = search_wikipedia("test", use_api=False)

        assert results == ["result"]
        mock_colbert.assert_called_once()


class TestDownloadBM25Index:
    """Tests for BM25 index download."""

    def test_download_with_default_cache(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Test downloading index with default cache directory."""
        calls: dict[str, Any] = {}

        def snapshot_download(**kwargs: Any) -> str:
            calls.update(kwargs)
            return str(tmp_path / "hf-cache")

        monkeypatch.setitem(
            sys.modules,
            "huggingface_hub",
            SimpleNamespace(snapshot_download=snapshot_download),
        )

        result = _download_bm25_index("test/repo")

        assert result == tmp_path / "hf-cache"
        assert calls["repo_id"] == "test/repo"
        assert calls["repo_type"] == "dataset"
        assert calls["local_dir_use_symlinks"] is False
        assert "local_dir" not in calls

    def test_download_with_custom_cache(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Test downloading index with custom cache directory."""
        calls: dict[str, Any] = {}

        def snapshot_download(**kwargs: Any) -> str:
            calls.update(kwargs)
            return str(tmp_path / "downloaded")

        monkeypatch.setitem(
            sys.modules,
            "huggingface_hub",
            SimpleNamespace(snapshot_download=snapshot_download),
        )

        target_dir = tmp_path / "bm25-cache"
        result = _download_bm25_index("test/repo", target_dir=target_dir)

        assert target_dir.exists()
        assert result == tmp_path / "downloaded"
        assert calls["repo_id"] == "test/repo"
        assert calls["repo_type"] == "dataset"
        assert calls["local_dir"] == str(target_dir)
        assert calls["local_dir_use_symlinks"] is False

    def test_download_missing_huggingface_hub(
        self, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Test error when huggingface_hub not installed."""
        real_import = builtins.__import__

        def import_side_effect(name: str, *args: Any, **kwargs: Any) -> Any:
            if name == "huggingface_hub":
                raise ImportError("missing")
            return real_import(name, *args, **kwargs)

        monkeypatch.setattr(builtins, "__import__", import_side_effect)

        with pytest.raises(ImportError, match="huggingface_hub not installed"):
            _download_bm25_index("test/repo")