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

"""Tests for ColBERT utility module."""

import copy
from collections.abc import Iterator
from unittest.mock import MagicMock, patch

import pytest

from opik_optimizer.utils.tools.colbert import (
    ColBERTv2,
    dotdict,
    colbertv2_get_request,
    colbertv2_post_request,
    _create_session_with_retries,
)


def _json_response(payload: dict) -> MagicMock:
    response = MagicMock()
    response.json.return_value = payload
    return response


@pytest.fixture
def mocked_requests_session() -> Iterator[MagicMock]:
    """
    Patch `requests.Session` construction inside the ColBERT module.

    Tests can then configure `.get` / `.post` return values or side effects without
    repeating the same patch boilerplate.
    """
    with patch("opik_optimizer.utils.tools.colbert.requests.Session") as session_cls:
        session = MagicMock()
        session_cls.return_value = session
        yield session


class TestDotDict:
    """Tests for dotdict helper class."""

    def test_attribute_access(self) -> None:
        d = dotdict({"name": "test", "value": 42})
        assert d.name == "test"
        assert d.value == 42

    def test_attribute_assignment(self) -> None:
        d = dotdict()
        d.name = "test"
        d.value = 42
        assert d["name"] == "test"
        assert d["value"] == 42

    def test_attribute_deletion(self) -> None:
        d = dotdict({"name": "test", "value": 42})
        del d.name
        assert "name" not in d
        assert d.value == 42

    def test_raises_attribute_error_for_missing_key(self) -> None:
        d = dotdict({"name": "test"})
        with pytest.raises(AttributeError, match="has no attribute"):
            _ = d.missing_key

    def test_dunder_attribute_passthrough(self) -> None:
        d = dotdict()
        # Should not raise - dunder attributes use parent behavior
        repr(d)
        str(d)

    def test_deepcopy(self) -> None:
        d = dotdict({"nested": {"value": 42}})
        copied = copy.deepcopy(d)
        assert copied.nested["value"] == 42
        # Modify original - copy should be independent
        d.nested["value"] = 100
        assert copied.nested["value"] == 42

    def test_dict_operations(self) -> None:
        d = dotdict({"a": 1, "b": 2})
        assert len(d) == 2
        assert list(d.keys()) == ["a", "b"]
        assert list(d.values()) == [1, 2]


class TestCreateSessionWithRetries:
    """Tests for _create_session_with_retries function."""

    def test_creates_session(self) -> None:
        session = _create_session_with_retries()
        assert session is not None
        # Check adapters are mounted
        assert "http://" in session.adapters
        assert "https://" in session.adapters

    def test_custom_max_retries(self) -> None:
        session = _create_session_with_retries(max_retries=10)
        assert session is not None


class TestColbertv2GetRequest:
    """Tests for colbertv2_get_request function."""

    def test_successful_request(self, mocked_requests_session: MagicMock) -> None:
        mock_response = {
            "topk": [
                {"text": "Result 1", "score": 0.9},
                {"text": "Result 2", "score": 0.8},
            ]
        }
        mocked_requests_session.get.return_value = _json_response(mock_response)

        results = colbertv2_get_request(
            url="http://localhost:8000/search",
            query="test query",
            k=2,
        )

        assert len(results) == 2
        assert results[0]["text"] == "Result 1"
        assert results[0]["long_text"] == "Result 1"

    def test_raises_on_k_greater_than_100(self) -> None:
        with pytest.raises(AssertionError, match="k <= 100"):
            colbertv2_get_request(
                url="http://localhost:8000/search",
                query="test query",
                k=101,
            )

    def test_handles_server_error_response(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response = {
            "error": True,
            "message": "Internal server error",
        }
        mocked_requests_session.get.return_value = _json_response(mock_response)

        with pytest.raises(Exception, match="ColBERTv2 server error"):
            colbertv2_get_request(
                url="http://localhost:8000/search",
                query="test query",
                k=5,
            )

    def test_retries_on_connection_error(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response_error = {
            "error": True,
            "message": "Cannot connect to host",
        }
        mock_response_success = {"topk": [{"text": "Result", "score": 0.9}]}
        mocked_requests_session.get.side_effect = [
            _json_response(mock_response_error),
            _json_response(mock_response_success),
        ]

        with patch("opik_optimizer.utils.tools.colbert.time.sleep"):
            results = colbertv2_get_request(
                url="http://localhost:8000/search",
                query="test query",
                k=1,
            )

        assert len(results) == 1

    def test_handles_unexpected_response_format(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response = {"unexpected_key": "value"}
        mocked_requests_session.get.return_value = _json_response(mock_response)

        with pytest.raises(Exception, match="Unexpected response format"):
            colbertv2_get_request(
                url="http://localhost:8000/search",
                query="test query",
                k=5,
            )

    def test_handles_request_exception(
        self, mocked_requests_session: MagicMock
    ) -> None:
        import requests

        mocked_requests_session.get.side_effect = requests.RequestException(
            "Network error"
        )

        with patch("opik_optimizer.utils.tools.colbert.time.sleep"):
            with pytest.raises(Exception, match="ColBERTv2 request failed"):
                colbertv2_get_request(
                    url="http://localhost:8000/search",
                    query="test query",
                    k=5,
                    max_retries=2,
                )


class TestColbertv2PostRequest:
    """Tests for colbertv2_post_request function."""

    def test_successful_post_request(self, mocked_requests_session: MagicMock) -> None:
        mock_response = {
            "topk": [
                {"text": "Result 1", "score": 0.9},
                {"text": "Result 2", "score": 0.8},
            ]
        }
        mocked_requests_session.post.return_value = _json_response(mock_response)

        results = colbertv2_post_request(
            url="http://localhost:8000/search",
            query="test query",
            k=2,
        )

        assert len(results) == 2
        assert results[0]["text"] == "Result 1"

    def test_handles_server_error_response(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response = {
            "error": True,
            "message": "Bad request",
        }
        mocked_requests_session.post.return_value = _json_response(mock_response)

        with pytest.raises(Exception, match="ColBERTv2 server error"):
            colbertv2_post_request(
                url="http://localhost:8000/search",
                query="test query",
                k=5,
            )

    def test_retries_on_connection_refused(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response_error = {
            "error": True,
            "message": "Connection refused",
        }
        mocked_requests_session.post.return_value = _json_response(mock_response_error)

        with patch("opik_optimizer.utils.tools.colbert.time.sleep"):
            with pytest.raises(Exception, match="ColBERTv2 server error"):
                colbertv2_post_request(
                    url="http://localhost:8000/search",
                    query="test query",
                    k=5,
                    max_retries=2,
                )


class TestColBERTv2Class:
    """Tests for ColBERTv2 wrapper class."""

    def test_initialization_with_defaults(self) -> None:
        client = ColBERTv2()
        assert client.url == "http://0.0.0.0"
        assert client.post_requests is False

    def test_initialization_with_port(self) -> None:
        client = ColBERTv2(url="http://localhost", port=8080)
        assert client.url == "http://localhost:8080"

    def test_initialization_with_post_requests(self) -> None:
        client = ColBERTv2(post_requests=True)
        assert client.post_requests is True

    def test_call_with_get_request(self, mocked_requests_session: MagicMock) -> None:
        mock_response = {
            "topk": [
                {"text": "Result 1", "score": 0.9},
                {"text": "Result 2", "score": 0.8},
            ]
        }
        mocked_requests_session.get.return_value = _json_response(mock_response)

        client = ColBERTv2(post_requests=False)
        results = client(query="test", k=2)

        assert len(results) == 2
        assert isinstance(results[0], dotdict)

    def test_call_with_post_request(self, mocked_requests_session: MagicMock) -> None:
        mock_response = {
            "topk": [
                {"text": "Result 1", "score": 0.9},
            ]
        }
        mocked_requests_session.post.return_value = _json_response(mock_response)

        client = ColBERTv2(post_requests=True)
        results = client(query="test", k=1)

        assert len(results) == 1
        mocked_requests_session.post.assert_called_once()

    def test_call_with_simplify_returns_strings(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response = {
            "topk": [
                {"text": "Result 1", "score": 0.9, "long_text": "Result 1 full"},
                {"text": "Result 2", "score": 0.8, "long_text": "Result 2 full"},
            ]
        }
        mocked_requests_session.get.return_value = _json_response(mock_response)

        client = ColBERTv2()
        results = client(query="test", k=2, simplify=True)

        assert len(results) == 2
        assert isinstance(results[0], str)

    def test_call_with_custom_max_retries(
        self, mocked_requests_session: MagicMock
    ) -> None:
        mock_response = {"topk": [{"text": "Result", "score": 0.9}]}
        mocked_requests_session.get.return_value = _json_response(mock_response)

        client = ColBERTv2()
        results = client(query="test", k=1, max_retries=10)

        assert len(results) == 1