dify/api/tests/integration_tests/vdb/opensearch/test_opensearch.py

from unittest.mock import MagicMock, patch

import pytest

from core.rag.datasource.vdb.field import Field
from core.rag.datasource.vdb.opensearch.opensearch_vector import OpenSearchConfig, OpenSearchVector
from core.rag.models.document import Document
from extensions import ext_redis


def get_example_text() -> str:
    return "This is a sample text for testing purposes."


@pytest.fixture(scope="module")
def setup_mock_redis():
    ext_redis.redis_client.get = MagicMock(return_value=None)
    ext_redis.redis_client.set = MagicMock(return_value=None)

    mock_redis_lock = MagicMock()
    mock_redis_lock.__enter__ = MagicMock()
    mock_redis_lock.__exit__ = MagicMock()
    ext_redis.redis_client.lock = MagicMock(return_value=mock_redis_lock)


class TestOpenSearchConfig:
    def test_to_opensearch_params(self):
        config = OpenSearchConfig(
            host="localhost",
            port=9200,
            secure=True,
            user="admin",
            password="password",
        )

        params = config.to_opensearch_params()

        assert params["hosts"] == [{"host": "localhost", "port": 9200}]
        assert params["use_ssl"] is True
        assert params["verify_certs"] is True
        assert params["connection_class"].__name__ == "Urllib3HttpConnection"
        assert params["http_auth"] == ("admin", "password")

    @patch("boto3.Session")
    @patch("core.rag.datasource.vdb.opensearch.opensearch_vector.Urllib3AWSV4SignerAuth")
    def test_to_opensearch_params_with_aws_managed_iam(
        self, mock_aws_signer_auth: MagicMock, mock_boto_session: MagicMock
    ):
        mock_credentials = MagicMock()
        mock_boto_session.return_value.get_credentials.return_value = mock_credentials

        mock_auth_instance = MagicMock()
        mock_aws_signer_auth.return_value = mock_auth_instance

        aws_region = "ap-southeast-2"
        aws_service = "aoss"
        host = f"aoss-endpoint.{aws_region}.aoss.amazonaws.com"
        port = 9201

        config = OpenSearchConfig(
            host=host,
            port=port,
            secure=True,
            auth_method="aws_managed_iam",
            aws_region=aws_region,
            aws_service=aws_service,
        )

        params = config.to_opensearch_params()

        assert params["hosts"] == [{"host": host, "port": port}]
        assert params["use_ssl"] is True
        assert params["verify_certs"] is True
        assert params["connection_class"].__name__ == "Urllib3HttpConnection"
        assert params["http_auth"] is mock_auth_instance

        mock_aws_signer_auth.assert_called_once_with(
            credentials=mock_credentials, region=aws_region, service=aws_service
        )
        assert mock_boto_session.return_value.get_credentials.called


class TestOpenSearchVector:
    def setup_method(self):
        self.collection_name = "test_collection"
        self.example_doc_id = "example_doc_id"
        self.vector = OpenSearchVector(
            collection_name=self.collection_name,
            config=OpenSearchConfig(host="localhost", port=9200, secure=False, user="admin", password="password"),
        )
        self.vector._client = MagicMock()

    @pytest.mark.parametrize(
        ("search_response", "expected_length", "expected_doc_id"),
        [
            (
                {
                    "hits": {
                        "total": {"value": 1},
                        "hits": [
                            {
                                "_source": {
                                    "page_content": get_example_text(),
                                    "metadata": {"document_id": "example_doc_id"},
                                }
                            }
                        ],
                    }
                },
                1,
                "example_doc_id",
            ),
            ({"hits": {"total": {"value": 0}, "hits": []}}, 0, None),
        ],
    )
    def test_search_by_full_text(self, search_response, expected_length, expected_doc_id):
        self.vector._client.search.return_value = search_response

        hits_by_full_text = self.vector.search_by_full_text(query=get_example_text())
        assert len(hits_by_full_text) == expected_length
        if expected_length > 0:
            assert hits_by_full_text[0].metadata["document_id"] == expected_doc_id

    def test_search_by_vector(self):
        vector = [0.1] * 128
        mock_response = {
            "hits": {
                "total": {"value": 1},
                "hits": [
                    {
                        "_source": {
                            Field.CONTENT_KEY.value: get_example_text(),
                            Field.METADATA_KEY.value: {"document_id": self.example_doc_id},
                        },
                        "_score": 1.0,
                    }
                ],
            }
        }
        self.vector._client.search.return_value = mock_response

        hits_by_vector = self.vector.search_by_vector(query_vector=vector)

        print("Hits by vector:", hits_by_vector)
        print("Expected document ID:", self.example_doc_id)
        print("Actual document ID:", hits_by_vector[0].metadata["document_id"] if hits_by_vector else "No hits")

        assert len(hits_by_vector) > 0, f"Expected at least one hit, got {len(hits_by_vector)}"
        assert hits_by_vector[0].metadata["document_id"] == self.example_doc_id, (
            f"Expected document ID {self.example_doc_id}, got {hits_by_vector[0].metadata['document_id']}"
        )

    def test_get_ids_by_metadata_field(self):
        mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}}
        self.vector._client.search.return_value = mock_response

        doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
        embedding = [0.1] * 128

        with patch("opensearchpy.helpers.bulk") as mock_bulk:
            mock_bulk.return_value = ([], [])
            self.vector.add_texts([doc], [embedding])

        ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
        assert len(ids) == 1
        assert ids[0] == "mock_id"

    def test_add_texts(self):
        self.vector._client.index.return_value = {"result": "created"}

        doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
        embedding = [0.1] * 128

        with patch("opensearchpy.helpers.bulk") as mock_bulk:
            mock_bulk.return_value = ([], [])
            self.vector.add_texts([doc], [embedding])

        mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}}
        self.vector._client.search.return_value = mock_response

        ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
        assert len(ids) == 1
        assert ids[0] == "mock_id"


@pytest.mark.usefixtures("setup_mock_redis")
class TestOpenSearchVectorWithRedis:
    def setup_method(self):
        self.tester = TestOpenSearchVector()

    def test_search_by_full_text(self):
        self.tester.setup_method()
        search_response = {
            "hits": {
                "total": {"value": 1},
                "hits": [
                    {"_source": {"page_content": get_example_text(), "metadata": {"document_id": "example_doc_id"}}}
                ],
            }
        }
        expected_length = 1
        expected_doc_id = "example_doc_id"
        self.tester.test_search_by_full_text(search_response, expected_length, expected_doc_id)

    def test_get_ids_by_metadata_field(self):
        self.tester.setup_method()
        self.tester.test_get_ids_by_metadata_field()

    def test_add_texts(self):
        self.tester.setup_method()
        self.tester.test_add_texts()

    def test_search_by_vector(self):
        self.tester.setup_method()
        self.tester.test_search_by_vector()