learningcircuit--local-deep-research/tests/utilities/test_es_utils.py

"""Tests for es_utils module."""

import tempfile
from pathlib import Path
from unittest.mock import Mock, patch

import pytest


class TestElasticsearchManagerInit:
    """Tests for ElasticsearchManager initialization."""

    def test_initializes_with_default_hosts(self):
        """Should initialize with default localhost host."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test-cluster",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ) as MockES:
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            ElasticsearchManager()

            MockES.assert_called_once()
            # Default hosts should be localhost
            call_args = MockES.call_args
            assert call_args[0][0] == ["http://localhost:9200"]

    def test_initializes_with_custom_hosts(self):
        """Should initialize with custom hosts."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ) as MockES:
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            ElasticsearchManager(hosts=["http://custom:9200"])

            call_args = MockES.call_args
            assert call_args[0][0] == ["http://custom:9200"]

    def test_initializes_with_basic_auth(self):
        """Should initialize with basic authentication."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ) as MockES:
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            ElasticsearchManager(username="user", password="pass")

            call_kwargs = MockES.call_args[1]
            assert call_kwargs["basic_auth"] == ("user", "pass")

    def test_initializes_with_api_key(self):
        """Should initialize with API key authentication."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ) as MockES:
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            ElasticsearchManager(api_key="my-api-key")

            call_kwargs = MockES.call_args[1]
            assert call_kwargs["api_key"] == "my-api-key"

    def test_initializes_with_cloud_id(self):
        """Should initialize with Elastic Cloud ID."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ) as MockES:
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            ElasticsearchManager(cloud_id="my-cloud-id")

            call_kwargs = MockES.call_args[1]
            assert call_kwargs["cloud_id"] == "my-cloud-id"

    def test_raises_connection_error_on_failure(self):
        """Should raise ConnectionError if connection fails."""
        mock_client = Mock()
        mock_client.info.side_effect = Exception("Connection refused")

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            with pytest.raises(ConnectionError):
                ElasticsearchManager()


class TestElasticsearchManagerCreateIndex:
    """Tests for ElasticsearchManager.create_index method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            mgr = ElasticsearchManager()
            return mgr

    def test_creates_index_successfully(self, manager):
        """Should create index and return True."""
        manager.client.indices.exists.return_value = False
        manager.client.indices.create.return_value = {"acknowledged": True}

        result = manager.create_index("test-index")

        assert result is True
        manager.client.indices.create.assert_called_once()

    def test_skips_existing_index(self, manager):
        """Should skip creation if index already exists."""
        manager.client.indices.exists.return_value = True

        result = manager.create_index("existing-index")

        assert result is True
        manager.client.indices.create.assert_not_called()

    def test_uses_custom_mappings(self, manager):
        """Should use provided custom mappings."""
        manager.client.indices.exists.return_value = False
        manager.client.indices.create.return_value = {"acknowledged": True}

        custom_mappings = {"properties": {"custom_field": {"type": "keyword"}}}
        manager.create_index("test-index", mappings=custom_mappings)

        call_kwargs = manager.client.indices.create.call_args[1]
        assert call_kwargs["mappings"] == custom_mappings

    def test_uses_custom_settings(self, manager):
        """Should use provided custom settings."""
        manager.client.indices.exists.return_value = False
        manager.client.indices.create.return_value = {"acknowledged": True}

        custom_settings = {"number_of_shards": 3}
        manager.create_index("test-index", settings=custom_settings)

        call_kwargs = manager.client.indices.create.call_args[1]
        assert call_kwargs["settings"] == custom_settings

    def test_returns_false_on_error(self, manager):
        """Should return False on error."""
        manager.client.indices.exists.return_value = False
        manager.client.indices.create.side_effect = Exception("Error")

        result = manager.create_index("test-index")

        assert result is False


class TestElasticsearchManagerDeleteIndex:
    """Tests for ElasticsearchManager.delete_index method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            return ElasticsearchManager()

    def test_deletes_index_successfully(self, manager):
        """Should delete index and return True."""
        manager.client.indices.exists.return_value = True
        manager.client.indices.delete.return_value = {"acknowledged": True}

        result = manager.delete_index("test-index")

        assert result is True
        manager.client.indices.delete.assert_called_once_with(
            index="test-index"
        )

    def test_skips_nonexistent_index(self, manager):
        """Should skip deletion if index doesn't exist."""
        manager.client.indices.exists.return_value = False

        result = manager.delete_index("nonexistent-index")

        assert result is True
        manager.client.indices.delete.assert_not_called()

    def test_returns_false_on_error(self, manager):
        """Should return False on error."""
        manager.client.indices.exists.return_value = True
        manager.client.indices.delete.side_effect = Exception("Error")

        result = manager.delete_index("test-index")

        assert result is False


class TestElasticsearchManagerIndexDocument:
    """Tests for ElasticsearchManager.index_document method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            return ElasticsearchManager()

    def test_indexes_document_successfully(self, manager):
        """Should index document and return ID."""
        manager.client.index.return_value = {"_id": "doc-123"}

        document = {"title": "Test", "content": "Content"}
        result = manager.index_document("test-index", document)

        assert result == "doc-123"
        manager.client.index.assert_called_once()

    def test_uses_provided_document_id(self, manager):
        """Should use provided document ID."""
        manager.client.index.return_value = {"_id": "custom-id"}

        document = {"title": "Test"}
        manager.index_document("test-index", document, document_id="custom-id")

        call_kwargs = manager.client.index.call_args[1]
        assert call_kwargs["id"] == "custom-id"

    def test_respects_refresh_parameter(self, manager):
        """Should pass refresh parameter to client."""
        manager.client.index.return_value = {"_id": "doc-123"}

        document = {"title": "Test"}
        manager.index_document("test-index", document, refresh=True)

        call_kwargs = manager.client.index.call_args[1]
        assert call_kwargs["refresh"] is True

    def test_returns_none_on_error(self, manager):
        """Should return None on error."""
        manager.client.index.side_effect = Exception("Error")

        document = {"title": "Test"}
        result = manager.index_document("test-index", document)

        assert result is None


class TestElasticsearchManagerBulkIndexDocuments:
    """Tests for ElasticsearchManager.bulk_index_documents method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            return ElasticsearchManager()

    def test_bulk_indexes_documents_successfully(self, manager):
        """Should bulk index documents and return count."""
        with patch(
            "local_deep_research.utilities.es_utils.bulk", return_value=(5, 0)
        ):
            documents = [{"title": f"Doc {i}"} for i in range(5)]
            result = manager.bulk_index_documents("test-index", documents)

            assert result == 5

    def test_uses_id_field_when_provided(self, manager):
        """Should use specified field as document ID."""
        with patch(
            "local_deep_research.utilities.es_utils.bulk", return_value=(2, 0)
        ) as mock_bulk:
            documents = [
                {"doc_id": "id-1", "title": "Doc 1"},
                {"doc_id": "id-2", "title": "Doc 2"},
            ]
            manager.bulk_index_documents(
                "test-index", documents, id_field="doc_id"
            )

            # Check that actions include _id
            call_args = mock_bulk.call_args
            actions = call_args[0][1]
            assert actions[0]["_id"] == "id-1"
            assert actions[1]["_id"] == "id-2"

    def test_returns_zero_on_error(self, manager):
        """Should return 0 on error."""
        with patch(
            "local_deep_research.utilities.es_utils.bulk",
            side_effect=Exception("Error"),
        ):
            documents = [{"title": "Doc"}]
            result = manager.bulk_index_documents("test-index", documents)

            assert result == 0


class TestElasticsearchManagerSearch:
    """Tests for ElasticsearchManager.search method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            return ElasticsearchManager()

    def test_searches_successfully(self, manager):
        """Should execute search and return response."""
        expected_response = {
            "hits": {
                "total": {"value": 1},
                "hits": [{"_source": {"title": "Test"}}],
            }
        }
        manager.client.search.return_value = expected_response

        result = manager.search("test-index", "test query")

        assert result == expected_response

    def test_uses_custom_fields(self, manager):
        """Should search in specified fields."""
        manager.client.search.return_value = {"hits": {"hits": []}}

        manager.search("test-index", "query", fields=["custom_field"])

        call_kwargs = manager.client.search.call_args[1]
        body = call_kwargs["body"]
        assert body["query"]["multi_match"]["fields"] == ["custom_field"]

    def test_respects_size_parameter(self, manager):
        """Should respect size parameter."""
        manager.client.search.return_value = {"hits": {"hits": []}}

        manager.search("test-index", "query", size=50)

        call_kwargs = manager.client.search.call_args[1]
        body = call_kwargs["body"]
        assert body["size"] == 50

    def test_includes_highlighting_by_default(self, manager):
        """Should include highlighting by default."""
        manager.client.search.return_value = {"hits": {"hits": []}}

        manager.search("test-index", "query")

        call_kwargs = manager.client.search.call_args[1]
        body = call_kwargs["body"]
        assert "highlight" in body

    def test_can_disable_highlighting(self, manager):
        """Should be able to disable highlighting."""
        manager.client.search.return_value = {"hits": {"hits": []}}

        manager.search("test-index", "query", highlight=False)

        call_kwargs = manager.client.search.call_args[1]
        body = call_kwargs["body"]
        assert "highlight" not in body

    def test_returns_error_dict_on_failure(self, manager):
        """Should return error dict on failure."""
        manager.client.search.side_effect = Exception("Search error")

        result = manager.search("test-index", "query")

        assert result == {"error": "Search failed"}


class TestElasticsearchManagerIndexFile:
    """Tests for ElasticsearchManager.index_file method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            return ElasticsearchManager()

    def test_returns_none_if_loader_not_available(self, manager):
        """Should return None if UnstructuredFileLoader not available."""
        # The import is lazy inside the function, so we patch at the source
        with patch.dict(
            "sys.modules",
            {"langchain_community.document_loaders": None},
        ):
            result = manager.index_file("test-index", "/path/to/file.txt")

            # Will raise ImportError, caught and returns None
            assert result is None

    def test_returns_none_on_error(self, manager):
        """Should return None on error."""
        # Force an error by using a non-existent file path without mocking loader
        result = manager.index_file("test-index", "/nonexistent/path/file.txt")

        assert result is None


class TestElasticsearchManagerIndexDirectory:
    """Tests for ElasticsearchManager.index_directory method."""

    @pytest.fixture
    def manager(self):
        """Create a manager with mocked client."""
        mock_client = Mock()
        mock_client.info.return_value = {
            "cluster_name": "test",
            "version": {"number": "8.0.0"},
        }

        with patch(
            "local_deep_research.utilities.es_utils.Elasticsearch",
            return_value=mock_client,
        ):
            from local_deep_research.utilities.es_utils import (
                ElasticsearchManager,
            )

            return ElasticsearchManager()

    def test_indexes_matching_files(self, manager):
        """Should index files matching patterns."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create test files
            Path(tmpdir, "test1.txt").write_text("Content 1")
            Path(tmpdir, "test2.txt").write_text("Content 2")
            Path(tmpdir, "test.pdf").write_text("PDF content")

            with patch.object(
                manager, "index_file", return_value="doc-id"
            ) as mock_index:
                result = manager.index_directory(
                    "test-index", tmpdir, file_patterns=["*.txt"]
                )

                assert result == 2
                assert mock_index.call_count == 2

    def test_returns_zero_on_error(self, manager):
        """Should return 0 on error."""
        result = manager.index_directory("test-index", "/nonexistent/path")

        assert result == 0

    def test_counts_successful_indexes_only(self, manager):
        """Should only count successfully indexed files."""
        with tempfile.TemporaryDirectory() as tmpdir:
            Path(tmpdir, "test1.txt").write_text("Content 1")
            Path(tmpdir, "test2.txt").write_text("Content 2")

            # First succeeds, second fails
            with patch.object(
                manager, "index_file", side_effect=["doc-1", None]
            ):
                result = manager.index_directory(
                    "test-index", tmpdir, file_patterns=["*.txt"]
                )

                assert result == 1