arc53--docsgpt/tests/vectorstore/test_elasticsearch.py

from unittest.mock import MagicMock, Mock, patch

import pytest


def _make_es_store(source_id="test-source"):
    """Helper to create an ElasticsearchStore with mocked deps."""
    # Reset class-level connection
    from application.vectorstore.elasticsearch import ElasticsearchStore

    ElasticsearchStore._es_connection = None

    with patch(
        "application.vectorstore.elasticsearch.settings"
    ) as mock_settings, patch.dict(
        "sys.modules", {"elasticsearch": MagicMock(), "elasticsearch.helpers": MagicMock()}
    ):
        mock_settings.ELASTIC_URL = "http://localhost:9200"
        mock_settings.ELASTIC_USERNAME = "elastic"
        mock_settings.ELASTIC_PASSWORD = "password"
        mock_settings.ELASTIC_CLOUD_ID = None
        mock_settings.ELASTIC_INDEX = "test_index"
        mock_settings.EMBEDDINGS_NAME = "test_model"

        import elasticsearch

        mock_es = MagicMock()
        elasticsearch.Elasticsearch.return_value = mock_es

        store = ElasticsearchStore(
            source_id=source_id,
            embeddings_key="key",
            index_name="test_index",
        )

        return store, mock_es, mock_settings


@pytest.mark.unit
class TestElasticsearchStoreInit:
    def test_source_id_cleaned(self):
        store, _, _ = _make_es_store(source_id="application/indexes/abc123/")
        assert store.source_id == "abc123"

    def test_init_with_url(self):
        store, mock_es, _ = _make_es_store()
        assert store.docsearch is mock_es
        assert store.index_name == "test_index"

    def test_init_with_cloud_id(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        ElasticsearchStore._es_connection = None

        with patch(
            "application.vectorstore.elasticsearch.settings"
        ) as mock_settings, patch.dict(
            "sys.modules", {"elasticsearch": MagicMock()}
        ):
            mock_settings.ELASTIC_URL = None
            mock_settings.ELASTIC_CLOUD_ID = "my-cloud-id"
            mock_settings.ELASTIC_USERNAME = "user"
            mock_settings.ELASTIC_PASSWORD = "pass"
            mock_settings.ELASTIC_INDEX = "idx"
            mock_settings.EMBEDDINGS_NAME = "model"

            store = ElasticsearchStore(
                source_id="src", embeddings_key="k", index_name="idx"
            )
            assert store.docsearch is not None

    def test_init_no_url_no_cloud_id_raises(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        ElasticsearchStore._es_connection = None

        with patch(
            "application.vectorstore.elasticsearch.settings"
        ) as mock_settings, patch.dict(
            "sys.modules", {"elasticsearch": MagicMock()}
        ):
            mock_settings.ELASTIC_URL = None
            mock_settings.ELASTIC_CLOUD_ID = None
            mock_settings.ELASTIC_INDEX = "idx"
            mock_settings.EMBEDDINGS_NAME = "model"

            with pytest.raises(ValueError, match="provide either"):
                ElasticsearchStore(source_id="src", embeddings_key="k")

    def test_reuses_class_connection(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        ElasticsearchStore._es_connection = None

        with patch(
            "application.vectorstore.elasticsearch.settings"
        ) as mock_settings, patch.dict(
            "sys.modules", {"elasticsearch": MagicMock()}
        ):
            mock_settings.ELASTIC_URL = "http://localhost:9200"
            mock_settings.ELASTIC_USERNAME = "user"
            mock_settings.ELASTIC_PASSWORD = "pass"
            mock_settings.ELASTIC_CLOUD_ID = None
            mock_settings.ELASTIC_INDEX = "idx"
            mock_settings.EMBEDDINGS_NAME = "model"

            import elasticsearch

            mock_es = MagicMock()
            elasticsearch.Elasticsearch.return_value = mock_es

            store1 = ElasticsearchStore(source_id="src1", embeddings_key="k")
            store2 = ElasticsearchStore(source_id="src2", embeddings_key="k")

            assert store1.docsearch is store2.docsearch
            elasticsearch.Elasticsearch.assert_called_once()


@pytest.mark.unit
class TestElasticsearchStoreSearch:
    def test_search_builds_query(self):
        store, mock_es, mock_settings = _make_es_store()

        mock_emb = Mock()
        mock_emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3])

        with patch.object(store, "_get_embeddings", return_value=mock_emb):
            mock_es.search.return_value = {
                "hits": {
                    "hits": [
                        {
                            "_source": {
                                "text": "doc1",
                                "metadata": {"source": "file.txt"},
                            }
                        },
                        {
                            "_source": {
                                "text": "doc2",
                                "metadata": {"source": "file2.txt"},
                            }
                        },
                    ]
                }
            }

            results = store.search("query", k=2)

        assert len(results) == 2
        assert results[0].page_content == "doc1"
        assert results[1].metadata == {"source": "file2.txt"}

    def test_search_empty_results(self):
        store, mock_es, _ = _make_es_store()

        mock_emb = Mock()
        mock_emb.embed_query = Mock(return_value=[0.1])

        with patch.object(store, "_get_embeddings", return_value=mock_emb):
            mock_es.search.return_value = {"hits": {"hits": []}}
            results = store.search("query")

        assert results == []


@pytest.mark.unit
class TestElasticsearchStoreAddTexts:
    def test_add_texts(self):
        store, mock_es, mock_settings = _make_es_store()

        mock_emb = Mock()
        mock_emb.embed_documents = Mock(return_value=[[0.1, 0.2], [0.3, 0.4]])

        mock_bulk = Mock(return_value=(2, 0))
        mock_helpers = MagicMock()
        mock_helpers.bulk = mock_bulk

        with patch.object(
            store, "_get_embeddings", return_value=mock_emb
        ), patch.object(
            store, "_create_index_if_not_exists"
        ), patch.dict(
            "sys.modules", {"elasticsearch.helpers": mock_helpers}
        ):
            ids = store.add_texts(
                ["text1", "text2"],
                metadatas=[{"a": 1}, {"b": 2}],
            )

        assert len(ids) == 2

    def test_add_texts_empty_raises(self):
        """Empty texts causes IndexError because code accesses vectors[0] unconditionally."""
        store, _, _ = _make_es_store()

        mock_emb = Mock()
        mock_emb.embed_documents = Mock(return_value=[])

        with patch.object(store, "_get_embeddings", return_value=mock_emb):
            with pytest.raises(IndexError):
                store.add_texts([], metadatas=[])


@pytest.mark.unit
class TestElasticsearchStoreDeleteIndex:
    def test_delete_index_calls_delete_by_query(self):
        store, mock_es, _ = _make_es_store(source_id="src1")

        store.delete_index()

        mock_es.delete_by_query.assert_called_once_with(
            index="test_index",
            query={"match": {"metadata.source_id.keyword": "src1"}},
        )


@pytest.mark.unit
class TestElasticsearchStoreIndex:
    def test_index_returns_mapping(self):
        store, _, _ = _make_es_store()

        mapping = store.index(dims_length=768)

        assert mapping["mappings"]["properties"]["vector"]["type"] == "dense_vector"
        assert mapping["mappings"]["properties"]["vector"]["dims"] == 768
        assert mapping["mappings"]["properties"]["vector"]["similarity"] == "cosine"

    def test_create_index_if_not_exists_existing(self):
        store, mock_es, _ = _make_es_store()
        mock_es.indices.exists.return_value = True

        store._create_index_if_not_exists("test_index", 768)

        mock_es.indices.create.assert_not_called()

    def test_create_index_if_not_exists_new(self):
        store, mock_es, _ = _make_es_store()
        mock_es.indices.exists.return_value = False

        store._create_index_if_not_exists("test_index", 768)

        mock_es.indices.create.assert_called_once()


@pytest.mark.unit
class TestElasticsearchStoreConnectToElasticsearch:
    def test_connect_with_url(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        with patch.dict("sys.modules", {"elasticsearch": MagicMock()}):
            import elasticsearch

            mock_es = MagicMock()
            elasticsearch.Elasticsearch.return_value = mock_es

            result = ElasticsearchStore.connect_to_elasticsearch(
                es_url="http://localhost:9200",
                username="user",
                password="pass",
            )
            assert result is mock_es

    def test_connect_with_both_raises(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        with patch.dict("sys.modules", {"elasticsearch": MagicMock()}):
            with pytest.raises(ValueError, match="Both es_url and cloud_id"):
                ElasticsearchStore.connect_to_elasticsearch(
                    es_url="http://localhost", cloud_id="cloud-123"
                )

    def test_connect_with_neither_raises(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        with patch.dict("sys.modules", {"elasticsearch": MagicMock()}):
            with pytest.raises(ValueError, match="provide either"):
                ElasticsearchStore.connect_to_elasticsearch()

    def test_connect_with_api_key(self):
        from application.vectorstore.elasticsearch import ElasticsearchStore

        with patch.dict("sys.modules", {"elasticsearch": MagicMock()}):
            import elasticsearch

            mock_es = MagicMock()
            elasticsearch.Elasticsearch.return_value = mock_es

            result = ElasticsearchStore.connect_to_elasticsearch(
                es_url="http://localhost:9200",
                api_key="my-api-key",
            )
            assert result is mock_es