"""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