alibaba--zvec/python/tests/test_embedding.py

# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import os
from http import HTTPStatus
from unittest.mock import MagicMock, patch, Mock

import numpy as np
import pytest
from zvec.extension import (
    BM25EmbeddingFunction,
    DefaultLocalDenseEmbedding,
    DefaultLocalSparseEmbedding,
    OpenAIDenseEmbedding,
    QwenDenseEmbedding,
    QwenSparseEmbedding,
)

# Environment variable to control integration tests
# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API/model tests
RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1"


# ----------------------------
# QwenDenseEmbedding Test Case
# ----------------------------
class TestQwenDenseEmbedding:
    def test_init_with_api_key(self):
        # Test initialization with explicit API key
        embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key")
        assert embedding_func.dimension == 128
        assert embedding_func.model == "text-embedding-v4"
        assert embedding_func._api_key == "test_key"

    @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "env_key"})
    def test_init_with_env_api_key(self):
        # Test initialization with API key from environment
        embedding_func = QwenDenseEmbedding(dimension=128)
        assert embedding_func._api_key == "env_key"

    @patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""})
    def test_init_with_empty_env_api_key(self):
        # Test initialization with empty API key from environment
        with pytest.raises(ValueError, match="DashScope API key is required"):
            QwenDenseEmbedding(dimension=128)

    def test_model_property(self):
        embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key")
        assert embedding_func.model == "text-embedding-v4"

        embedding_func = QwenDenseEmbedding(
            dimension=128, model="custom-model", api_key="test_key"
        )
        assert embedding_func.model == "custom-model"

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_with_empty_text(self, mock_require_module):
        # Test embed method with empty text raises ValueError
        embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key")

        with pytest.raises(
            ValueError, match="Input text cannot be empty or whitespace only"
        ):
            embedding_func.embed("")

        with pytest.raises(TypeError):
            embedding_func.embed(None)

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_success(self, mock_require_module):
        # Test successful embedding
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        mock_response.output = {"embeddings": [{"embedding": [0.1, 0.2, 0.3]}]}
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenDenseEmbedding(dimension=3, api_key="test_key")
        # Clear cache to avoid interference
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test text")

        assert result == [0.1, 0.2, 0.3]
        mock_dashscope.TextEmbedding.call.assert_called_once_with(
            model="text-embedding-v4",
            input="test text",
            dimension=3,
            output_type="dense",
        )

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_http_error(self, mock_require_module):
        # Test embedding with HTTP error
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.BAD_REQUEST
        mock_response.message = "Bad Request"
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key")
        embedding_func.embed.cache_clear()

        with pytest.raises(ValueError):
            embedding_func.embed("test text")

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_invalid_response(self, mock_require_module):
        # Test embedding with invalid response (wrong number of embeddings)
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        mock_response.output = {"embeddings": []}
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key")
        embedding_func.embed.cache_clear()

        with pytest.raises(ValueError):
            embedding_func.embed("test text")

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    def test_real_embed_success(self):
        """Integration test with real DashScope API.

        To run this test, set environment variable:
            export ZVEC_RUN_INTEGRATION_TESTS=1
            export DASHSCOPE_API_KEY=your-api-key
        """
        embedding_func = QwenDenseEmbedding(dimension=128)
        dense = embedding_func("test text")
        assert len(dense) == 128


# ----------------------------
# QwenSparseEmbedding Test Case
# ----------------------------
class TestQwenSparseEmbedding:
    """Test suite for QwenSparseEmbedding (Qwen sparse embedding via DashScope API)."""

    def test_init_with_api_key(self):
        """Test initialization with explicit API key."""
        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        assert embedding_func._dimension == 1024
        assert embedding_func.model == "text-embedding-v4"
        assert embedding_func._api_key == "test_key"
        # encoding_type defaults to "query" via extra_params
        assert embedding_func.extra_params.get("encoding_type", "query") == "query"

    def test_init_with_custom_encoding_type(self):
        """Test initialization with custom encoding type."""
        embedding_func = QwenSparseEmbedding(
            dimension=1024, encoding_type="document", api_key="test_key"
        )
        assert embedding_func.extra_params.get("encoding_type") == "document"

    @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "env_key"})
    def test_init_with_env_api_key(self):
        """Test initialization with API key from environment."""
        embedding_func = QwenSparseEmbedding(dimension=1024)
        assert embedding_func._api_key == "env_key"

    @patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""})
    def test_init_without_api_key(self):
        """Test initialization fails without API key."""
        with pytest.raises(ValueError, match="DashScope API key is required"):
            QwenSparseEmbedding(dimension=1024)

    def test_model_property(self):
        """Test model property."""
        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        assert embedding_func.model == "text-embedding-v4"

        embedding_func = QwenSparseEmbedding(
            dimension=1024, model="text-embedding-v3", api_key="test_key"
        )
        assert embedding_func.model == "text-embedding-v3"

    def test_encoding_type_property(self):
        """Test encoding_type via extra_params."""
        query_emb = QwenSparseEmbedding(
            dimension=1024, encoding_type="query", api_key="test_key"
        )
        assert query_emb.extra_params.get("encoding_type") == "query"

        doc_emb = QwenSparseEmbedding(
            dimension=1024, encoding_type="document", api_key="test_key"
        )
        assert doc_emb.extra_params.get("encoding_type") == "document"

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_with_empty_text(self, mock_require_module):
        """Test embed method with empty text raises ValueError."""
        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")

        with pytest.raises(
            ValueError, match="Input text cannot be empty or whitespace only"
        ):
            embedding_func.embed("")

        with pytest.raises(
            ValueError, match="Input text cannot be empty or whitespace only"
        ):
            embedding_func.embed("   ")

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_with_non_string_input(self, mock_require_module):
        """Test embed method with non-string input raises TypeError."""
        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            embedding_func.embed(123)

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            embedding_func.embed(None)

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_success(self, mock_require_module):
        """Test successful sparse embedding generation."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        # Sparse embedding returns array of {index, value, token} objects
        mock_response.output = {
            "embeddings": [
                {
                    "sparse_embedding": [
                        {"index": 10, "value": 0.5, "token": "机器"},
                        {"index": 245, "value": 0.8, "token": "学习"},
                        {"index": 1023, "value": 1.2, "token": "算法"},
                    ]
                }
            ]
        }
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        # Clear cache to avoid interference
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test text")

        # Verify result is a dict
        assert isinstance(result, dict)
        # Verify keys are integers
        assert all(isinstance(k, int) for k in result.keys())
        # Verify values are floats
        assert all(isinstance(v, float) for v in result.values())
        # Verify all values are positive
        assert all(v > 0 for v in result.values())
        # Verify sorted by indices
        keys = list(result.keys())
        assert keys == sorted(keys)
        # Verify specific keys
        assert keys == [10, 245, 1023]

        mock_dashscope.TextEmbedding.call.assert_called_once_with(
            model="text-embedding-v4",
            input="test text",
            dimension=1024,
            output_type="sparse",
            text_type="query",
        )

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_with_document_encoding_type(self, mock_require_module):
        """Test embedding with document encoding type."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        mock_response.output = {
            "embeddings": [
                {
                    "sparse_embedding": [
                        {"index": 5, "value": 0.3, "token": "文档"},
                        {"index": 100, "value": 0.7, "token": "内容"},
                        {"index": 500, "value": 0.9, "token": "检索"},
                    ]
                }
            ]
        }
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(
            dimension=1024, encoding_type="document", api_key="test_key"
        )
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test document")

        assert isinstance(result, dict)
        assert list(result.keys()) == [5, 100, 500]

        # Verify text_type parameter is "document"
        call_args = mock_dashscope.TextEmbedding.call.call_args
        assert call_args[1]["text_type"] == "document"
        assert call_args[1]["output_type"] == "sparse"

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_output_sorted_by_indices(self, mock_require_module):
        """Test that output is always sorted by indices in ascending order."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        # Return unsorted indices
        mock_response.output = {
            "embeddings": [
                {
                    "sparse_embedding": [
                        {"index": 9999, "value": 1.5, "token": "A"},
                        {"index": 5, "value": 2.0, "token": "B"},
                        {"index": 1234, "value": 0.8, "token": "C"},
                        {"index": 77, "value": 3.2, "token": "D"},
                        {"index": 500, "value": 1.1, "token": "E"},
                    ]
                }
            ]
        }
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test sorting")

        # Verify keys are sorted
        result_keys = list(result.keys())
        assert result_keys == sorted(result_keys)
        # Verify expected sorted order
        assert result_keys == [5, 77, 500, 1234, 9999]

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_filters_zero_values(self, mock_require_module):
        """Test that zero and negative values are filtered out."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        # Include zero and negative values
        mock_response.output = {
            "embeddings": [
                {
                    "sparse_embedding": [
                        {"index": 10, "value": 0.5, "token": "正"},
                        {
                            "index": 20,
                            "value": 0.0,
                            "token": "零",
                        },  # Should be filtered
                        {
                            "index": 30,
                            "value": -0.3,
                            "token": "负",
                        },  # Should be filtered
                        {"index": 40, "value": 0.8, "token": "正"},
                        {
                            "index": 50,
                            "value": 0.0,
                            "token": "零",
                        },  # Should be filtered
                    ]
                }
            ]
        }
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test filtering")

        # Only positive values should remain
        assert list(result.keys()) == [10, 40]
        assert all(v > 0 for v in result.values())

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_http_error(self, mock_require_module):
        """Test embedding with HTTP error."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.BAD_REQUEST
        mock_response.message = "Bad Request"
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()

        with pytest.raises(ValueError, match="DashScope API error"):
            embedding_func.embed("test text")

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_invalid_response_no_embeddings(self, mock_require_module):
        """Test embedding with invalid response (no embeddings)."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        mock_response.output = {"embeddings": []}
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()

        with pytest.raises(ValueError, match="Expected exactly 1 embedding"):
            embedding_func.embed("test text")

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_invalid_response_not_dict(self, mock_require_module):
        """Test embedding with invalid response (sparse_embedding not list)."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        # sparse_embedding should be list, not dict
        mock_response.output = {
            "embeddings": [{"sparse_embedding": {"index": 10, "value": 0.5}}]
        }
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()

        with pytest.raises(
            ValueError, match="'sparse_embedding' field is missing or not a list"
        ):
            embedding_func.embed("test text")

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_callable_interface(self, mock_require_module):
        """Test that embedding function is callable."""
        mock_dashscope = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = HTTPStatus.OK
        mock_response.output = {
            "embeddings": [
                {
                    "sparse_embedding": [
                        {"index": 100, "value": 1.0, "token": "测试"},
                        {"index": 200, "value": 0.5, "token": "调用"},
                    ]
                }
            ]
        }
        mock_dashscope.TextEmbedding.call.return_value = mock_response
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()

        # Test calling the function directly
        result = embedding_func("test text")
        assert isinstance(result, dict)
        assert list(result.keys()) == [100, 200]

    @patch("zvec.extension.qwen_function.require_module")
    def test_embed_api_connection_error(self, mock_require_module):
        """Test handling of API connection errors."""
        mock_dashscope = MagicMock()
        mock_dashscope.TextEmbedding.call.side_effect = Exception("Connection timeout")
        mock_require_module.return_value = mock_dashscope

        embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key")
        embedding_func.embed.cache_clear()

        with pytest.raises(RuntimeError, match="Failed to call DashScope API"):
            embedding_func.embed("test text")

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    def test_real_embed_success(self):
        """Integration test with real DashScope API.

        To run this test, set environment variable:
            export ZVEC_RUN_INTEGRATION_TESTS=1
            export DASHSCOPE_API_KEY=your-api-key
        """
        # Test query embedding
        query_emb = QwenSparseEmbedding(dimension=1024, encoding_type="query")
        query_vec = query_emb.embed("machine learning")

        assert isinstance(query_vec, dict)
        assert len(query_vec) > 0
        assert all(isinstance(k, int) for k in query_vec.keys())
        assert all(isinstance(v, float) and v > 0 for v in query_vec.values())

        # Verify sorted output
        keys = list(query_vec.keys())
        assert keys == sorted(keys)

        # Test document embedding
        doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type="document")
        doc_vec = doc_emb.embed("Machine learning is a subset of AI")

        assert isinstance(doc_vec, dict)
        assert len(doc_vec) > 0

        # Verify sorted output
        doc_keys = list(doc_vec.keys())
        assert doc_keys == sorted(doc_keys)


# ----------------------------
# OpenAIDenseEmbedding Test Case
# ----------------------------
class TestOpenAIDenseEmbedding:
    def test_init_with_api_key(self):
        """Test initialization with explicit API key."""
        embedding_func = OpenAIDenseEmbedding(api_key="sk-test-key")
        assert embedding_func.dimension == 1536  # Default for text-embedding-3-small
        assert embedding_func.model == "text-embedding-3-small"
        assert embedding_func._api_key == "sk-test-key"

    @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-env-key"})
    def test_init_with_env_api_key(self):
        """Test initialization with API key from environment."""
        embedding_func = OpenAIDenseEmbedding()
        assert embedding_func._api_key == "sk-env-key"

    @patch.dict(os.environ, {"OPENAI_API_KEY": ""})
    def test_init_without_api_key(self):
        """Test initialization fails without API key."""
        with pytest.raises(ValueError, match="OpenAI API key is required"):
            OpenAIDenseEmbedding()

    def test_init_with_custom_dimension(self):
        """Test initialization with custom dimension."""
        embedding_func = OpenAIDenseEmbedding(
            model="text-embedding-3-large", dimension=1024, api_key="sk-test"
        )
        assert embedding_func.dimension == 1024
        assert embedding_func.model == "text-embedding-3-large"

    def test_init_with_base_url(self):
        """Test initialization with custom base URL."""
        embedding_func = OpenAIDenseEmbedding(
            api_key="sk-test", base_url="https://custom.openai.com/"
        )
        assert embedding_func._base_url == "https://custom.openai.com/"

    def test_model_property(self):
        """Test model property."""
        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")
        assert embedding_func.model == "text-embedding-3-small"

        embedding_func = OpenAIDenseEmbedding(
            model="text-embedding-ada-002", api_key="sk-test"
        )
        assert embedding_func.model == "text-embedding-ada-002"

    def test_extra_params(self):
        """Test extra_params property."""
        # Test without extra params
        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")
        assert embedding_func.extra_params == {}

        # Test with extra params
        embedding_func = OpenAIDenseEmbedding(
            api_key="sk-test",
            encoding_format="float",
            user="test-user",
        )
        assert embedding_func.extra_params == {
            "encoding_format": "float",
            "user": "test-user",
        }

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_with_empty_text(self, mock_require_module):
        """Test embed method with empty text raises ValueError."""
        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")

        with pytest.raises(
            ValueError, match="Input text cannot be empty or whitespace only"
        ):
            embedding_func.embed("")

        with pytest.raises(
            ValueError, match="Input text cannot be empty or whitespace only"
        ):
            embedding_func.embed("   ")

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_with_non_string_input(self, mock_require_module):
        """Test embed method with non-string input raises TypeError."""
        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            embedding_func.embed(123)

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            embedding_func.embed(None)

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_success(self, mock_require_module):
        """Test successful embedding generation."""
        # Mock OpenAI client
        mock_openai = Mock()
        mock_client = Mock()
        mock_response = Mock()

        # Create mock embedding data
        fake_embedding = [0.1, 0.2, 0.3]
        mock_embedding_obj = Mock()
        mock_embedding_obj.embedding = fake_embedding
        mock_response.data = [mock_embedding_obj]

        mock_client.embeddings.create.return_value = mock_response
        mock_openai.OpenAI.return_value = mock_client
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(dimension=3, api_key="sk-test")
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test text")

        assert result == [0.1, 0.2, 0.3]
        mock_client.embeddings.create.assert_called_once_with(
            model="text-embedding-3-small", input="test text", dimensions=3
        )

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_with_custom_model(self, mock_require_module):
        """Test embedding with custom model."""
        mock_openai = Mock()
        mock_client = Mock()
        mock_response = Mock()

        fake_embedding = [0.1] * 1536
        mock_embedding_obj = Mock()
        mock_embedding_obj.embedding = fake_embedding
        mock_response.data = [mock_embedding_obj]

        mock_client.embeddings.create.return_value = mock_response
        mock_openai.OpenAI.return_value = mock_client
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(
            model="text-embedding-ada-002", api_key="sk-test"
        )
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test text")

        assert len(result) == 1536
        mock_client.embeddings.create.assert_called_once_with(
            model="text-embedding-ada-002", input="test text"
        )

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_api_error(self, mock_require_module):
        """Test handling of API errors."""
        mock_openai = Mock()
        mock_client = Mock()

        # Simulate API error
        api_error = Mock()
        api_error.__class__.__name__ = "APIError"
        mock_openai.APIError = type("APIError", (Exception,), {})
        mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {})

        mock_client.embeddings.create.side_effect = mock_openai.APIError(
            "Rate limit exceeded"
        )
        mock_openai.OpenAI.return_value = mock_client
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")
        embedding_func.embed.cache_clear()

        with pytest.raises(RuntimeError, match="Failed to call OpenAI API"):
            embedding_func.embed("test text")

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_invalid_response(self, mock_require_module):
        """Test handling of invalid API response."""
        mock_openai = Mock()
        mock_client = Mock()
        mock_response = Mock()

        # Empty response data
        mock_response.data = []

        mock_client.embeddings.create.return_value = mock_response
        mock_openai.OpenAI.return_value = mock_client
        mock_openai.APIError = type("APIError", (Exception,), {})
        mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {})
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")
        embedding_func.embed.cache_clear()

        with pytest.raises(ValueError, match="no embedding data returned"):
            embedding_func.embed("test text")

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_dimension_mismatch(self, mock_require_module):
        """Test handling of dimension mismatch."""
        mock_openai = Mock()
        mock_client = Mock()
        mock_response = Mock()

        # Return embedding with wrong dimension
        fake_embedding = [0.1] * 512
        mock_embedding_obj = Mock()
        mock_embedding_obj.embedding = fake_embedding
        mock_response.data = [mock_embedding_obj]

        mock_client.embeddings.create.return_value = mock_response
        mock_openai.OpenAI.return_value = mock_client
        mock_openai.APIError = type("APIError", (Exception,), {})
        mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {})
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(dimension=1536, api_key="sk-test")
        embedding_func.embed.cache_clear()

        with pytest.raises(ValueError, match="Dimension mismatch"):
            embedding_func.embed("test text")

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_callable(self, mock_require_module):
        """Test that embedding function is callable."""
        mock_openai = Mock()
        mock_client = Mock()
        mock_response = Mock()

        fake_embedding = [0.1] * 1536
        mock_embedding_obj = Mock()
        mock_embedding_obj.embedding = fake_embedding
        mock_response.data = [mock_embedding_obj]

        mock_client.embeddings.create.return_value = mock_response
        mock_openai.OpenAI.return_value = mock_client
        mock_openai.APIError = type("APIError", (Exception,), {})
        mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {})
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(api_key="sk-test")
        embedding_func.embed.cache_clear()

        # Test calling the function directly
        result = embedding_func("test text")
        assert isinstance(result, list)
        assert len(result) == 1536

    @patch("zvec.extension.openai_function.require_module")
    def test_embed_with_base_url(self, mock_require_module):
        """Test embedding with custom base URL."""
        mock_openai = Mock()
        mock_client = Mock()
        mock_response = Mock()

        fake_embedding = [0.1] * 1536
        mock_embedding_obj = Mock()
        mock_embedding_obj.embedding = fake_embedding
        mock_response.data = [mock_embedding_obj]

        mock_client.embeddings.create.return_value = mock_response
        mock_openai.OpenAI.return_value = mock_client
        mock_openai.APIError = type("APIError", (Exception,), {})
        mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {})
        mock_require_module.return_value = mock_openai

        embedding_func = OpenAIDenseEmbedding(
            api_key="sk-test", base_url="https://custom.openai.com/"
        )
        embedding_func.embed.cache_clear()
        result = embedding_func.embed("test text")

        # Verify client was created with custom base URL
        mock_openai.OpenAI.assert_called_once_with(
            api_key="sk-test", base_url="https://custom.openai.com/"
        )
        assert len(result) == 1536

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    def test_real_embed_success(self):
        """Integration test with real OpenAI API.

        To run this test, set environment variable:
            export ZVEC_RUN_INTEGRATION_TESTS=1
            export OPENAI_API_KEY=sk-...
        """
        embedding_func = OpenAIDenseEmbedding(
            model="text-embedding-v4",
            dimension=256,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        vector = embedding_func.embed("Hello, world!")
        assert len(vector) == 256
        assert isinstance(vector, list)
        assert all(isinstance(x, float) for x in vector)


# ----------------------------
# DefaultLocalDenseEmbedding Test Case
# ----------------------------
class TestDefaultLocalDenseEmbedding:
    """Test cases for DefaultLocalDenseEmbedding."""

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_init_success(self, mock_require_module):
        """Test successful initialization with mocked model."""
        # Mock sentence_transformers module
        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_model.device = "cpu"
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        # Initialize embedding function
        emb_func = DefaultLocalDenseEmbedding()

        # Assertions
        assert emb_func.dimension == 384
        assert emb_func.model_name == "all-MiniLM-L6-v2"
        assert emb_func.model_source == "huggingface"
        assert emb_func.device == "cpu"
        mock_st.SentenceTransformer.assert_called_once_with(
            "all-MiniLM-L6-v2", device=None, trust_remote_code=True
        )

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_init_with_custom_device(self, mock_require_module):
        """Test initialization with custom device."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_model.device = "cuda"
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding(device="cuda")

        assert emb_func.device == "cuda"
        mock_st.SentenceTransformer.assert_called_once_with(
            "all-MiniLM-L6-v2", device="cuda", trust_remote_code=True
        )

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_init_with_modelscope(self, mock_require_module):
        """Test initialization with ModelScope as model source."""
        mock_st = Mock()
        mock_ms = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_model.device = "cpu"
        mock_st.SentenceTransformer.return_value = mock_model

        def require_module_side_effect(module_name):
            if module_name == "sentence_transformers":
                return mock_st
            elif module_name == "modelscope":
                return mock_ms
            raise ImportError(f"No module named '{module_name}'")

        mock_require_module.side_effect = require_module_side_effect

        # Mock snapshot_download at the correct import location
        with patch(
            "modelscope.hub.snapshot_download.snapshot_download",
            return_value="/path/to/cached/model",
        ):
            emb_func = DefaultLocalDenseEmbedding(model_source="modelscope")

        # Assertions
        assert emb_func.dimension == 384
        assert emb_func.model_name == "iic/nlp_gte_sentence-embedding_chinese-small"
        assert emb_func.model_source == "modelscope"

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_init_with_invalid_model_source(self, mock_require_module):
        """Test initialization with invalid model_source raises ValueError."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        with pytest.raises(ValueError, match="Invalid model_source"):
            DefaultLocalDenseEmbedding(model_source="invalid_source")

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_success(self, mock_require_module):
        """Test successful embedding generation."""
        # Mock embedding output
        fake_embedding = np.random.rand(384).astype(np.float32)

        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384

        # Configure encode method
        mock_model.encode = Mock(return_value=fake_embedding)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding()
        result = emb_func.embed("Hello, world!")

        # Assertions
        assert isinstance(result, list)
        assert len(result) == 384
        assert all(isinstance(x, float) for x in result)
        mock_model.encode.assert_called_once_with(
            "Hello, world!",
            convert_to_numpy=True,
            normalize_embeddings=True,
            batch_size=32,
        )

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_with_normalization(self, mock_require_module):
        """Test embedding with L2 normalization."""
        # Create a normalized vector
        fake_embedding = np.random.rand(384).astype(np.float32)
        fake_embedding = fake_embedding / np.linalg.norm(fake_embedding)

        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384

        # Configure encode method
        mock_model.encode = Mock(return_value=fake_embedding)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding(normalize_embeddings=True)
        result = emb_func.embed("Test sentence")

        # Check if vector is normalized (L2 norm should be close to 1.0)
        result_array = np.array(result)
        norm = np.linalg.norm(result_array)
        assert abs(norm - 1.0) < 1e-5

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_empty_string(self, mock_require_module):
        """Test embedding with empty string raises ValueError."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding()

        with pytest.raises(ValueError, match="Input text cannot be empty"):
            emb_func.embed("")

        with pytest.raises(ValueError, match="Input text cannot be empty"):
            emb_func.embed("   ")

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_non_string_input(self, mock_require_module):
        """Test embedding with non-string input raises TypeError."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding()

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            emb_func.embed(123)

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            emb_func.embed(None)

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_callable(self, mock_require_module):
        """Test that embedding function is callable."""
        fake_embedding = np.random.rand(384).astype(np.float32)

        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384

        # Configure encode method
        mock_model.encode = Mock(return_value=fake_embedding)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding()

        # Test calling the function directly
        result = emb_func("Test text")
        assert isinstance(result, list)
        assert len(result) == 384

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_semantic_similarity(self, mock_require_module):
        """Test semantic similarity between similar and different texts."""
        # Create mock embeddings for similar and different texts
        similar_emb_1 = np.array([1.0, 0.0, 0.0] + [0.0] * 381, dtype=np.float32)
        similar_emb_2 = np.array([0.9, 0.1, 0.0] + [0.0] * 381, dtype=np.float32)
        different_emb = np.array([0.0, 0.0, 1.0] + [0.0] * 381, dtype=np.float32)

        # Normalize
        similar_emb_1 = similar_emb_1 / np.linalg.norm(similar_emb_1)
        similar_emb_2 = similar_emb_2 / np.linalg.norm(similar_emb_2)
        different_emb = different_emb / np.linalg.norm(different_emb)

        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384

        # Configure encode method with side_effect for multiple calls
        mock_model.encode = Mock(
            side_effect=[similar_emb_1, similar_emb_2, different_emb]
        )

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding()

        v1 = emb_func.embed("The cat sits on the mat")
        v2 = emb_func.embed("A feline rests on a rug")
        v3 = emb_func.embed("Python programming")

        # Calculate similarities
        similarity_high = np.dot(v1, v2)
        similarity_low = np.dot(v1, v3)

        assert similarity_high > similarity_low

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_model_loading_error(self, mock_require_module):
        """Test handling of model loading failure."""
        # Clear model cache
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()
        mock_st = Mock()
        mock_st.SentenceTransformer.side_effect = Exception("Model not found")
        mock_require_module.return_value = mock_st

        with pytest.raises(
            ValueError, match="Failed to load Sentence Transformer model"
        ):
            DefaultLocalDenseEmbedding()

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_modelscope_import_error(self, mock_require_module):
        """Test handling of ModelScope import error."""
        mock_st = Mock()

        def require_module_side_effect(module_name):
            if module_name == "sentence_transformers":
                return mock_st
            elif module_name == "modelscope":
                raise ImportError("No module named 'modelscope'")

        mock_require_module.side_effect = require_module_side_effect

        with pytest.raises(
            ImportError, match="ModelScope support requires the 'modelscope' package"
        ):
            DefaultLocalDenseEmbedding(model_source="modelscope")

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_dimension_mismatch(self, mock_require_module):
        """Test handling of dimension mismatch in embedding output."""
        # Return embedding with wrong dimension
        fake_embedding = np.random.rand(256).astype(np.float32)

        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384

        # Configure encode method
        mock_model.encode = Mock(return_value=fake_embedding)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        emb_func = DefaultLocalDenseEmbedding()

        with pytest.raises(ValueError, match="Dimension mismatch"):
            emb_func.embed("Test text")

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    def test_real_embedding_generation(self):
        """Integration test with real model (requires sentence-transformers).

        To run this test, set environment variable:
            export ZVEC_RUN_INTEGRATION_TESTS=1

        Note: First run will download the model (~80MB).
        """
        emb_func = DefaultLocalDenseEmbedding()

        # Test basic embedding
        vector = emb_func.embed("Hello, world!")
        assert len(vector) == 384
        assert isinstance(vector, list)
        assert all(isinstance(x, float) for x in vector)

        # Test normalization
        norm = np.linalg.norm(vector)
        assert abs(norm - 1.0) < 1e-5

        # Test semantic similarity
        v1 = emb_func.embed("The cat sits on the mat")
        v2 = emb_func.embed("A feline rests on a rug")
        v3 = emb_func.embed("Python programming language")

        similarity_high = np.dot(v1, v2)
        similarity_low = np.dot(v1, v3)
        assert similarity_high > similarity_low

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_model_properties(self, mock_require_module):
        """Test model_name and model_source properties."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.get_sentence_embedding_dimension.return_value = 384
        mock_model.device = "cpu"
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        # Test Hugging Face
        emb_func_hf = DefaultLocalDenseEmbedding(model_source="huggingface")
        assert emb_func_hf.model_name == "all-MiniLM-L6-v2"
        assert emb_func_hf.model_source == "huggingface"

        # Test ModelScope
        with patch(
            "modelscope.hub.snapshot_download.snapshot_download",
            return_value="/path/to/model",
        ):
            mock_ms = Mock()
            mock_require_module.side_effect = lambda m: (
                mock_st if m == "sentence_transformers" else mock_ms
            )
            emb_func_ms = DefaultLocalDenseEmbedding(model_source="modelscope")
            assert (
                emb_func_ms.model_name == "iic/nlp_gte_sentence-embedding_chinese-small"
            )
            assert emb_func_ms.model_source == "modelscope"


# -----------------------------------
# DefaultLocalSparseEmbedding Test Case
# -----------------------------------
class TestDefaultLocalSparseEmbedding:
    """Test suite for DefaultLocalSparseEmbedding (SPLADE sparse embedding).

    Note:
        DefaultLocalSparseEmbedding uses naver/splade-cocondenser-ensembledistil
        instead of naver/splade-v3 because:

        - splade-v3 is a gated model requiring Hugging Face authentication
        - cocondenser-ensembledistil is publicly accessible
        - Performance difference is minimal (~2%)
        - Avoids "Access to model is restricted" errors

        This allows all users to run tests without authentication setup.
    """

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_init_success(self, mock_require_module):
        """Test successful initialization.

        Verifies that DefaultLocalSparseEmbedding initializes with the publicly
        accessible naver/splade-cocondenser-ensembledistil model instead of
        the gated naver/splade-v3 model.
        """
        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()

        assert sparse_emb.model_name == "naver/splade-cocondenser-ensembledistil"
        assert sparse_emb.model_source == "huggingface"
        assert sparse_emb.device == "cpu"
        mock_st.SentenceTransformer.assert_called_once_with(
            "naver/splade-cocondenser-ensembledistil",
            device=None,
            trust_remote_code=True,
        )

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_init_with_custom_device(self, mock_require_module):
        """Test initialization with custom device."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cuda"
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding(device="cuda")

        assert sparse_emb.device == "cuda"
        mock_st.SentenceTransformer.assert_called_once_with(
            "naver/splade-cocondenser-ensembledistil",
            device="cuda",
            trust_remote_code=True,
        )

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_success(self, mock_require_module):
        """Test successful sparse embedding generation with official API."""
        import numpy as np

        # Clear model cache to ensure fresh mock
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        # Create a mock sparse matrix that simulates scipy.sparse behavior
        # The code will call: sparse_matrix[0].toarray().flatten()
        mock_sparse_matrix = Mock()

        # Create a dense array representation with vocab_size=30522
        vocab_size = 30522
        dense_array = np.zeros(vocab_size)
        # Set specific non-zero values at indices [10, 245, 1023, 5678]
        dense_array[10] = 0.5
        dense_array[245] = 0.8
        dense_array[1023] = 1.2
        dense_array[5678] = 0.3

        # Mock the method chain: sparse_matrix[0].toarray().flatten()
        mock_row = Mock()
        mock_dense = Mock()
        mock_row.toarray.return_value = mock_dense
        mock_dense.flatten.return_value = dense_array
        mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)

        # Also mock hasattr check for 'toarray'
        mock_sparse_matrix.toarray = Mock()

        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"

        # Configure mock methods to return sparse matrix
        # Must set return_value BEFORE hasattr() check in the code
        mock_model.encode_query = Mock(return_value=mock_sparse_matrix)
        mock_model.encode_document = Mock(return_value=mock_sparse_matrix)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()
        result = sparse_emb.embed("machine learning")

        # Verify result is a dictionary
        assert isinstance(result, dict)
        # Verify keys are integers and values are floats
        assert all(isinstance(k, int) for k in result.keys())
        assert all(isinstance(v, float) for v in result.values())
        # Verify all values are positive
        assert all(v > 0 for v in result.values())
        # Sparse vectors should have specific dimensions
        assert len(result) == 4

        # Verify output is sorted by indices (keys)
        keys = list(result.keys())
        assert keys == sorted(keys), (
            "Sparse vector keys must be sorted in ascending order"
        )

        # Verify expected keys
        assert keys == [10, 245, 1023, 5678]

        # Verify encode_query was called with a list
        mock_model.encode_query.assert_called_once()
        call_args = mock_model.encode_query.call_args[0][0]
        assert isinstance(call_args, list)
        assert call_args == ["machine learning"]

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_empty_input(self, mock_require_module):
        """Test embedding with empty input."""
        mock_st = Mock()
        mock_model = Mock()
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()

        with pytest.raises(ValueError, match="Input text cannot be empty"):
            sparse_emb.embed("")

        with pytest.raises(ValueError, match="Input text cannot be empty"):
            sparse_emb.embed("   ")

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_embed_non_string_input(self, mock_require_module):
        """Test embedding with non-string input."""
        mock_st = Mock()
        mock_model = Mock()
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            sparse_emb.embed(123)

        with pytest.raises(TypeError, match="Expected 'input' to be str"):
            sparse_emb.embed(["text"])

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_callable_interface(self, mock_require_module):
        """Test that DefaultSparseEmbedding is callable."""
        import numpy as np

        # Clear model cache
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        # Create a mock sparse matrix
        mock_sparse_matrix = Mock()

        # Create a dense array representation with vocab_size=30522
        vocab_size = 30522
        dense_array = np.zeros(vocab_size)
        # Set specific non-zero values at indices [100, 200, 300]
        dense_array[100] = 1.0
        dense_array[200] = 0.5
        dense_array[300] = 0.8

        # Mock the method chain: sparse_matrix[0].toarray().flatten()
        mock_row = Mock()
        mock_dense = Mock()
        mock_row.toarray.return_value = mock_dense
        mock_dense.flatten.return_value = dense_array
        mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)

        # Also mock hasattr check for 'toarray'
        mock_sparse_matrix.toarray = Mock()

        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"

        # Configure mock methods
        mock_model.encode_query = Mock(return_value=mock_sparse_matrix)
        mock_model.encode_document = Mock(return_value=mock_sparse_matrix)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()

        # Test callable interface
        result = sparse_emb("test input")
        assert isinstance(result, dict)
        assert all(isinstance(k, int) for k in result.keys())

        # Verify sorted output
        keys = list(result.keys())
        assert keys == sorted(keys), "Callable interface must also return sorted keys"
        assert keys == [100, 200, 300]

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_model_loading_failure(self, mock_require_module):
        """Test handling of model loading failure."""
        # Clear model cache to ensure the test actually tries to load the model
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        mock_st = Mock()
        mock_st.SentenceTransformer.side_effect = Exception("Model not found")
        mock_require_module.return_value = mock_st

        with pytest.raises(
            ValueError, match="Failed to load Sentence Transformer model"
        ):
            DefaultLocalSparseEmbedding()

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_inference_failure(self, mock_require_module):
        """Test handling of inference failure."""
        # Clear model cache
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"

        # Configure mock methods to raise RuntimeError
        mock_model.encode_query = Mock(side_effect=RuntimeError("CUDA out of memory"))
        mock_model.encode_document = Mock(
            side_effect=RuntimeError("CUDA out of memory")
        )

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()

        with pytest.raises(RuntimeError, match="Failed to generate sparse embedding"):
            sparse_emb.embed("test input")

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_sparse_vector_properties(self, mock_require_module):
        """Test properties of sparse vectors (sparsity, non-zero values, sorted order)."""
        import numpy as np

        # Clear model cache
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        # Create a mock sparse matrix that simulates scipy.sparse behavior
        # The code will call: sparse_matrix[0].toarray().flatten()
        mock_sparse_matrix = Mock()

        # Create a dense array representation with vocab_size=30522
        vocab_size = 30522
        dense_array = np.zeros(vocab_size)
        # Set specific non-zero values at indices [50, 100, 200, 400, 500]
        dense_array[50] = 3.0
        dense_array[100] = 2.0
        dense_array[200] = 1.5
        dense_array[400] = 2.5
        dense_array[500] = 1.8

        # Mock the method chain: sparse_matrix[0].toarray().flatten()
        mock_row = Mock()
        mock_dense = Mock()
        mock_row.toarray.return_value = mock_dense
        mock_dense.flatten.return_value = dense_array
        mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)

        # Also mock hasattr check for 'toarray'
        mock_sparse_matrix.toarray = Mock()

        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"

        # Configure mock methods
        mock_model.encode_query = Mock(return_value=mock_sparse_matrix)
        mock_model.encode_document = Mock(return_value=mock_sparse_matrix)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()
        result = sparse_emb.embed("test")

        # Verify sparsity: result should have much fewer dimensions than vocab_size
        assert len(result) < vocab_size
        # All values should be positive
        assert all(v > 0 for v in result.values())

        # Verify keys are sorted in ascending order
        keys = list(result.keys())
        assert keys == sorted(keys), "Sparse vector keys must be sorted"

        # Verify the specific non-zero indices are present and sorted
        # Expected order: [50, 100, 200, 400, 500] (sorted)
        expected_keys = [50, 100, 200, 400, 500]
        assert keys == expected_keys, f"Expected {expected_keys}, got {keys}"

        # First key should be smallest
        if len(result) > 0:
            first_key = next(iter(result.keys()))
            assert first_key == min(result.keys()), "First key must be the smallest"

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_output_sorted_by_indices(self, mock_require_module):
        """Test that output dictionary is always sorted by indices (keys) in ascending order."""
        import numpy as np

        # Clear model cache
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        # Create sparse output with deliberately out-of-order indices
        # Non-sequential indices: 9999, 5, 1234, 77, 500
        mock_sparse_matrix = Mock()

        # Create a dense array representation with vocab_size=30522
        vocab_size = 30522
        dense_array = np.zeros(vocab_size)
        # Set specific non-zero values at out-of-order indices
        dense_array[9999] = 1.5
        dense_array[5] = 2.0
        dense_array[1234] = 0.8
        dense_array[77] = 3.2
        dense_array[500] = 1.1

        # Mock the method chain: sparse_matrix[0].toarray().flatten()
        mock_row = Mock()
        mock_dense = Mock()
        mock_row.toarray.return_value = mock_dense
        mock_dense.flatten.return_value = dense_array
        mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)

        # Also mock hasattr check for 'toarray'
        mock_sparse_matrix.toarray = Mock()

        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"

        # Configure mock methods
        mock_model.encode_query = Mock(return_value=mock_sparse_matrix)
        mock_model.encode_document = Mock(return_value=mock_sparse_matrix)

        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding()
        result = sparse_emb.embed("test sorting")

        # Extract keys from result
        result_keys = list(result.keys())

        # Verify keys are sorted
        assert result_keys == sorted(result_keys), (
            f"Keys must be sorted in ascending order. "
            f"Got: {result_keys}, Expected: {sorted(result_keys)}"
        )

        # Verify expected keys are present and in correct order
        # Expected sorted order: [5, 77, 500, 1234, 9999]
        expected_sorted_keys = [5, 77, 500, 1234, 9999]
        assert result_keys == expected_sorted_keys, (
            f"All expected keys should be present in sorted order. "
            f"Expected: {expected_sorted_keys}, Got: {result_keys}"
        )

        # Verify first and last keys
        assert result_keys[0] == 5, "First key must be minimum"
        assert result_keys[-1] == 9999, "Last key must be maximum"

        # Verify iteration order matches sorted order
        for i, (key, value) in enumerate(result.items()):
            if i > 0:
                prev_key = list(result.keys())[i - 1]
                assert key > prev_key, (
                    f"Key at position {i} must be greater than previous key"
                )

    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_device_property(self, mock_require_module):
        """Test device property returns correct device."""
        mock_st = Mock()
        mock_model = Mock()
        mock_model.device = "cuda"
        mock_st.SentenceTransformer.return_value = mock_model
        mock_require_module.return_value = mock_st

        sparse_emb = DefaultLocalSparseEmbedding(device="cuda")
        assert sparse_emb.device == "cuda"

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download",
    )
    @patch("zvec.extension.sentence_transformer_function.require_module")
    def test_modelscope_source(self, mock_require_module):
        """Test initialization with ModelScope source."""
        mock_st = Mock()
        mock_ms = Mock()
        mock_model = Mock()
        mock_model.device = "cpu"
        mock_st.SentenceTransformer.return_value = mock_model

        # Mock ModelScope snapshot_download
        with patch(
            "modelscope.hub.snapshot_download.snapshot_download",
            return_value="/cache/splade-cocondenser",
        ):
            mock_require_module.side_effect = lambda m: (
                mock_st if m == "sentence_transformers" else mock_ms
            )

            sparse_emb = DefaultLocalSparseEmbedding(model_source="modelscope")

            assert sparse_emb.model_name == "naver/splade-cocondenser-ensembledistil"
            assert sparse_emb.model_source == "modelscope"

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download",
    )
    def test_integration_real_model(self):
        """Integration test with real SPLADE model (requires model download).

        This test uses naver/splade-cocondenser-ensembledistil instead of
        naver/splade-v3 because splade-v3 requires Hugging Face authentication.
        The cocondenser-ensembledistil model is publicly accessible and provides
        comparable performance.

        To run this test:
            export ZVEC_RUN_INTEGRATION_TESTS=1
            pytest tests/test_embedding.py::TestDefaultSparseEmbedding::test_integration_real_model -v

        Note: First run will download ~100MB model from Hugging Face.

        Alternative models:
            If you have access to splade-v3, you can create a custom embedding
            class following the example in DefaultSparseEmbedding docstring.
        """
        # Clear model cache to ensure fresh load
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        sparse_emb = DefaultLocalSparseEmbedding()

        # Test with real input
        text = "machine learning and artificial intelligence"
        result = sparse_emb.embed(text)

        # Verify result structure
        assert isinstance(result, dict)
        assert len(result) > 0
        assert all(isinstance(k, int) and k >= 0 for k in result.keys())
        assert all(isinstance(v, float) and v > 0 for v in result.values())

        # SPLADE typically produces 100-300 non-zero dimensions
        assert 50 < len(result) < 500

        # Verify keys are sorted in ascending order
        keys = list(result.keys())
        assert keys == sorted(keys), "Real model output must be sorted by indices"

        # Test callable interface
        result2 = sparse_emb(text)
        assert result == result2

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1",
    )
    def test_integration_multiple_inputs(self):
        """Integration test with multiple different inputs."""
        # Clear model cache
        from zvec.extension.sentence_transformer_embedding_function import (
            DefaultLocalSparseEmbedding,
        )

        DefaultLocalSparseEmbedding.clear_cache()

        sparse_emb = DefaultLocalSparseEmbedding()

        texts = [
            "Hello, world!",
            "Machine learning is fascinating",
            "Python programming language",
        ]

        results = [sparse_emb.embed(text) for text in texts]

        # All results should be different
        assert len(results) == 3
        assert all(isinstance(r, dict) for r in results)

        # Different inputs should produce different sparse vectors
        assert results[0] != results[1]
        assert results[1] != results[2]

        # All results must be sorted by indices
        for i, result in enumerate(results):
            keys = list(result.keys())
            assert keys == sorted(keys), f"Result {i} must have sorted keys"


# ----------------------------
# BM25EmbeddingFunction Test Case
# ----------------------------
class TestBM25EmbeddingFunction:
    """Test suite for BM25EmbeddingFunction (BM25-based sparse embedding using DashText SDK)."""

    def test_init_with_built_in_encoder(self):
        """Test successful initialization with built-in encoder (no corpus)."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            # Test with default language (Chinese)
            bm25 = BM25EmbeddingFunction()

            assert bm25.corpus_size == 0
            assert bm25.encoding_type == "query"
            assert bm25.language == "zh"
            mock_dashtext.SparseVectorEncoder.default.assert_called_once_with(name="zh")

    def test_init_with_custom_encoder(self):
        """Test successful initialization with custom encoder (with corpus)."""
        corpus = [
            "a cat is a feline and likes to purr",
            "a dog is the human's best friend",
            "a bird is a beautiful animal that can fly",
        ]

        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction(corpus=corpus, b=0.75, k1=1.2)

            assert bm25.corpus_size == 3
            assert bm25.encoding_type == "query"
            mock_dashtext.SparseVectorEncoder.assert_called_once_with(b=0.75, k1=1.2)
            mock_encoder.train.assert_called_once_with(corpus)

    def test_init_with_empty_corpus(self):
        """Test initialization with empty corpus raises ValueError."""
        with pytest.raises(ValueError, match="Corpus must be a non-empty list"):
            BM25EmbeddingFunction(corpus=[])

    def test_init_with_invalid_corpus(self):
        """Test initialization with invalid corpus elements."""
        with pytest.raises(ValueError, match="All corpus documents must be strings"):
            BM25EmbeddingFunction(corpus=["text", 123, "another"])

        with pytest.raises(ValueError, match="All corpus documents must be strings"):
            BM25EmbeddingFunction(corpus=[None, "text"])

    def test_init_with_language_parameter(self):
        """Test initialization with different language settings."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            # Test English language
            bm25_en = BM25EmbeddingFunction(language="en")
            assert bm25_en.language == "en"
            mock_dashtext.SparseVectorEncoder.default.assert_called_with(name="en")

    def test_init_with_encoding_type(self):
        """Test initialization with different encoding types."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            # Test document encoding type
            bm25_doc = BM25EmbeddingFunction(encoding_type="document")
            assert bm25_doc.encoding_type == "document"

    def test_init_with_missing_dashtext_library(self):
        """Test initialization fails when dashtext library is not installed."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_require.side_effect = ImportError("dashtext package is required")

            with pytest.raises(ImportError, match="dashtext package is required"):
                BM25EmbeddingFunction()

    def test_embed_with_query_encoding(self):
        """Test successful sparse embedding generation with query encoding."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()

            # Mock encode_queries to return sparse vector
            mock_encoder.encode_queries.return_value = {
                5: 0.89,
                12: 1.45,
                23: 0.67,
                45: 1.12,
            }

            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction(encoding_type="query")
            # Clear LRU cache to ensure fresh call
            bm25.embed.cache_clear()
            result = bm25.embed("cat purr loud")

            # Verify result structure
            assert isinstance(result, dict)
            assert all(isinstance(k, int) for k in result.keys())
            assert all(isinstance(v, float) for v in result.values())

            # Verify all values are positive
            assert all(v > 0 for v in result.values())

            # Verify output is sorted by indices
            keys = list(result.keys())
            assert keys == sorted(keys), "Output must be sorted by indices"

            # Verify expected keys from mock response
            assert result == {5: 0.89, 12: 1.45, 23: 0.67, 45: 1.12}

            # Verify encode_queries was called
            mock_encoder.encode_queries.assert_called_once_with("cat purr loud")

    def test_embed_with_document_encoding(self):
        """Test successful sparse embedding generation with document encoding."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()

            # Mock encode_documents to return sparse vector
            mock_encoder.encode_documents.return_value = {10: 1.5, 20: 2.3}

            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction(encoding_type="document")
            bm25.embed.cache_clear()
            result = bm25.embed("document text")

            assert result == {10: 1.5, 20: 2.3}
            mock_encoder.encode_documents.assert_called_once_with("document text")

    def test_embed_with_empty_input(self):
        """Test embedding with empty input raises ValueError."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction()

            with pytest.raises(ValueError, match="Input text cannot be empty"):
                bm25.embed("")

            with pytest.raises(ValueError, match="Input text cannot be empty"):
                bm25.embed("   ")

    def test_embed_with_non_string_input(self):
        """Test embedding with non-string input raises TypeError."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction()

            # Test with hashable non-string types - should get our custom error message
            with pytest.raises(TypeError, match="Expected 'input' to be str"):
                bm25.embed(123)

            with pytest.raises(TypeError, match="Expected 'input' to be str"):
                bm25.embed(None)

            # Test with unhashable type (list)
            # Note: lru_cache raises TypeError("unhashable type: 'list'") before our type check
            # This is still a valid type error, just caught at a different layer
            with pytest.raises(TypeError, match="unhashable type"):
                bm25.embed(["text"])

    def test_embed_callable_interface(self):
        """Test that BM25EmbeddingFunction is callable."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_encoder.encode_queries.return_value = {10: 1.5}
            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction()
            bm25.embed.cache_clear()

            # Test callable interface
            result = bm25("test query")
            assert isinstance(result, dict)
            assert 10 in result

    def test_embed_output_sorted_by_indices(self):
        """Test that output is always sorted by indices in ascending order."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()

            # Mock encode_queries with unsorted indices
            mock_encoder.encode_queries.return_value = {
                9999: 1.5,
                5: 2.0,
                1234: 0.8,
                77: 3.2,
                500: 1.1,
            }

            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction()
            bm25.embed.cache_clear()
            result = bm25.embed("test query")

            # Verify keys are sorted
            result_keys = list(result.keys())
            assert result_keys == sorted(result_keys), (
                f"Keys must be sorted. Got: {result_keys}, Expected: {sorted(result_keys)}"
            )

            # Verify expected sorted order: [5, 77, 500, 1234, 9999]
            expected_keys = [5, 77, 500, 1234, 9999]
            assert result_keys == expected_keys

    def test_embed_filters_zero_values(self):
        """Test that zero and negative values are filtered out."""
        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()

            # Mock encode_queries with zero and negative values
            mock_encoder.encode_queries.return_value = {
                0: 1.5,  # Positive - should be included
                1: 0.0,  # Zero - should be filtered
                2: -0.5,  # Negative - should be filtered
            }

            mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction()
            bm25.embed.cache_clear()
            result = bm25.embed("test")

            # Only positive token should be in result
            assert 0 in result
            assert 1 not in result  # Zero value filtered
            assert 2 not in result  # Negative value filtered
            assert all(v > 0 for v in result.values())

    def test_properties(self):
        """Test property accessors."""
        corpus = ["doc1", "doc2", "doc3"]

        with patch(
            "zvec.extension.bm25_embedding_function.require_module"
        ) as mock_require:
            mock_dashtext = Mock()
            mock_encoder = Mock()
            mock_dashtext.SparseVectorEncoder.return_value = mock_encoder
            mock_require.return_value = mock_dashtext

            bm25 = BM25EmbeddingFunction(
                corpus=corpus,
                encoding_type="document",
                language="en",
                b=0.8,
                k1=1.5,
                custom_param="test",
            )

            assert bm25.corpus_size == 3
            assert bm25.encoding_type == "document"
            assert bm25.language == "en"
            assert bm25.extra_params == {"custom_param": "test"}

    @pytest.mark.skipif(
        not RUN_INTEGRATION_TESTS,
        reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
    )
    def test_real_dashtext_bm25_embedding(self):
        """Integration test with real DashText library.

        To run this test:
            export ZVEC_RUN_INTEGRATION_TESTS=1
            pip install dashtext

        Note: This test requires the dashtext package to be installed.
        """
        # Test built-in encoder (Chinese)
        bm25_zh = BM25EmbeddingFunction(language="zh", encoding_type="query")

        query_zh = "什么是向量检索服务"
        result_zh = bm25_zh.embed(query_zh)

        assert isinstance(result_zh, dict)
        assert len(result_zh) > 0
        assert all(isinstance(k, int) for k in result_zh.keys())
        assert all(isinstance(v, float) and v > 0 for v in result_zh.values())

        # Verify sorted output
        keys = list(result_zh.keys())
        assert keys == sorted(keys), "Real DashText BM25 output must be sorted"

        # Test custom corpus
        corpus = [
            "The cat sits on the mat",
            "The dog plays in the garden",
            "Birds fly in the sky",
            "Fish swim in the water",
        ]

        bm25_custom = BM25EmbeddingFunction(corpus=corpus, encoding_type="query")

        query_en = "cat on mat"
        result_en = bm25_custom.embed(query_en)

        assert isinstance(result_en, dict)
        assert len(result_en) > 0
        assert all(isinstance(k, int) for k in result_en.keys())
        assert all(isinstance(v, float) and v > 0 for v in result_en.values())

        # Test callable interface
        result2 = bm25_custom(query_en)
        assert result_en == result2

        # Verify properties
        assert bm25_custom.corpus_size == 4