test calls to index doc create delete

Author: monneyboi

poliloom/poliloom/cli.py

@@ -716,22 +716,12 @@ def garbage_collect():
                 f"📋 Previous dump timestamp: {previous_dump.last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC"
             )
 
-            # Clean up missing entities
+            # Clean up missing entities (also removes from search index)
             click.echo("⏳ Cleaning up entities using two-dump validation...")
-            deleted_entity_ids = CurrentImportEntity.cleanup_missing(
+            deleted_entity_count = CurrentImportEntity.cleanup_missing(
                 session, previous_dump.last_modified
             )
-            click.echo(f"  • Soft-deleted {len(deleted_entity_ids)} entities")
-
-            # Remove deleted entities from search index
-            if deleted_entity_ids:
-                from poliloom.search import SearchService
-
-                search_service = SearchService()
-                search_service.delete_documents(deleted_entity_ids)
-                click.echo(
-                    f"  • Removed {len(deleted_entity_ids)} entities from search index"
-                )
+            click.echo(f"  • Soft-deleted {deleted_entity_count} entities")
 
             # Clean up missing statements
             click.echo("⏳ Cleaning up statements using two-dump validation...")
@@ -746,7 +736,7 @@ def garbage_collect():
             )
 
             total_deleted = (
-                len(deleted_entity_ids)
+                deleted_entity_count
                 + statement_counts["properties_marked_deleted"]
                 + statement_counts["relations_marked_deleted"]
             )

poliloom/poliloom/models/wikidata.py

@@ -824,20 +824,26 @@ class CurrentImportEntity(Base):
 
     @classmethod
     def cleanup_missing(
-        cls, session: Session, previous_dump_timestamp: datetime
-    ) -> list[str]:
+        cls,
+        session: Session,
+        previous_dump_timestamp: datetime,
+    ) -> int:
         """
         Soft-delete entities using two-dump validation strategy.
         Only deletes entities missing from current dump AND older than previous dump.
         This prevents race conditions from incorrectly deleting recently added entities.
 
+        Also removes deleted entities from the search index.
+
         Args:
             session: Database session
             previous_dump_timestamp: Last modified timestamp of the previous dump.
 
         Returns:
-            List of wikidata_ids that were soft-deleted
+            Number of entities that were soft-deleted
         """
+        from poliloom.search import SearchService
+
         # Only delete if: NOT in current dump AND older than previous dump
         # Convert timezone-aware timestamp to naive for database comparison
         previous_dump_naive = previous_dump_timestamp.replace(tzinfo=None)
@@ -855,7 +861,14 @@ def cleanup_missing(
             {"previous_dump_timestamp": previous_dump_naive},
         )
 
-        return [row[0] for row in deleted_result.fetchall()]
+        deleted_ids = [row[0] for row in deleted_result.fetchall()]
+
+        # Remove deleted entities from search index
+        if deleted_ids:
+            search_service = SearchService()
+            search_service.delete_documents(deleted_ids)
+
+        return len(deleted_ids)
 
     @classmethod
     def clear_tracking_table(cls, session: Session) -> None:

poliloom/tests/models/test_wikidata.py

@@ -1,5 +1,7 @@
 """Tests for WikidataEntity model."""
 
+from unittest.mock import patch, Mock
+
 from sqlalchemy.dialects.postgresql import insert
 
 from poliloom.models import WikidataEntity, WikidataRelation, RelationType
@@ -856,3 +858,151 @@ def test_entity_with_subclass_relation_to_hierarchy(self, db_session):
         # Position should remain
         count = db_session.execute(text("SELECT COUNT(*) FROM positions")).scalar()
         assert count == 1
+
+
+class TestCleanupOutsideHierarchySearchService:
+    """Test that cleanup_outside_hierarchy calls search service for deletion.
+
+    Uses Location because it has SearchIndexedMixin (Position does not).
+    """
+
+    def _create_hierarchy(self, db_session, root_id, child_ids):
+        """Helper to create a hierarchy with root and children."""
+        # Create wikidata entities
+        entities = [{"wikidata_id": root_id, "name": f"Root {root_id}"}]
+        for child_id in child_ids:
+            entities.append({"wikidata_id": child_id, "name": f"Child {child_id}"})
+
+        stmt = insert(WikidataEntity).values(entities)
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+
+        # Create subclass relations
+        relations = []
+        for child_id in child_ids:
+            relations.append(
+                {
+                    "parent_entity_id": root_id,
+                    "child_entity_id": child_id,
+                    "relation_type": RelationType.SUBCLASS_OF,
+                    "statement_id": f"{child_id}$subclass-of-{root_id}",
+                }
+            )
+
+        if relations:
+            stmt = insert(WikidataRelation).values(relations)
+            stmt = stmt.on_conflict_do_nothing(index_elements=["statement_id"])
+            db_session.execute(stmt)
+
+        db_session.flush()
+
+    def _create_location_in_hierarchy(self, db_session, location_id, class_id):
+        """Helper to create a location that is an instance of a hierarchy class."""
+        from poliloom.models import Location
+
+        # Create wikidata entity for the location
+        stmt = insert(WikidataEntity).values(
+            [{"wikidata_id": location_id, "name": f"Location {location_id}"}]
+        )
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+
+        # Create location record
+        stmt = insert(Location.__table__).values([{"wikidata_id": location_id}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+
+        # Create instance_of relation to the class
+        stmt = insert(WikidataRelation).values(
+            [
+                {
+                    "parent_entity_id": class_id,
+                    "child_entity_id": location_id,
+                    "relation_type": RelationType.INSTANCE_OF,
+                    "statement_id": f"{location_id}$instance-of-{class_id}",
+                }
+            ]
+        )
+        stmt = stmt.on_conflict_do_nothing(index_elements=["statement_id"])
+        db_session.execute(stmt)
+        db_session.flush()
+
+    def _create_orphan_location(self, db_session, location_id):
+        """Helper to create a location with no hierarchy relations."""
+        from poliloom.models import Location
+
+        # Create wikidata entity
+        stmt = insert(WikidataEntity).values(
+            [{"wikidata_id": location_id, "name": f"Orphan Location {location_id}"}]
+        )
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+
+        # Create location record (no relations)
+        stmt = insert(Location.__table__).values([{"wikidata_id": location_id}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+        db_session.flush()
+
+    def test_cleanup_calls_delete_documents(self, db_session):
+        """Test that cleanup_outside_hierarchy calls delete_documents on search service."""
+        from poliloom.models import Location
+
+        # Create hierarchy (Q486972 is "human settlement" in Location._hierarchy_roots)
+        self._create_hierarchy(db_session, "Q486972", ["Q100"])
+
+        # Create a valid location
+        self._create_location_in_hierarchy(db_session, "Q200", "Q100")
+
+        # Create orphan locations that will be deleted
+        self._create_orphan_location(db_session, "Q300")
+        self._create_orphan_location(db_session, "Q301")
+
+        # Mock the SearchService (imported locally in cleanup_outside_hierarchy)
+        mock_search_service = Mock()
+        with patch("poliloom.search.SearchService", return_value=mock_search_service):
+            stats = Location.cleanup_outside_hierarchy(db_session, dry_run=False)
+
+        # Verify delete_documents was called with the orphan IDs
+        assert stats["entities_removed"] == 2
+        mock_search_service.delete_documents.assert_called_once()
+        deleted_ids = mock_search_service.delete_documents.call_args[0][0]
+        assert set(deleted_ids) == {"Q300", "Q301"}
+
+    def test_cleanup_does_not_call_delete_when_nothing_removed(self, db_session):
+        """Test that delete_documents is not called when no entities are removed."""
+        from poliloom.models import Location
+
+        # Create hierarchy
+        self._create_hierarchy(db_session, "Q486972", ["Q100"])
+
+        # Create only valid locations (no orphans)
+        self._create_location_in_hierarchy(db_session, "Q200", "Q100")
+
+        # Mock the SearchService
+        mock_search_service = Mock()
+        with patch("poliloom.search.SearchService", return_value=mock_search_service):
+            stats = Location.cleanup_outside_hierarchy(db_session, dry_run=False)
+
+        # No entities removed, so delete_documents should not be called
+        assert stats["entities_removed"] == 0
+        mock_search_service.delete_documents.assert_not_called()
+
+    def test_dry_run_does_not_call_delete_documents(self, db_session):
+        """Test that dry_run=True does not call delete_documents."""
+        from poliloom.models import Location
+
+        # Create hierarchy
+        self._create_hierarchy(db_session, "Q486972", ["Q100"])
+
+        # Create orphan location
+        self._create_orphan_location(db_session, "Q300")
+
+        # Mock the SearchService
+        mock_search_service = Mock()
+        with patch("poliloom.search.SearchService", return_value=mock_search_service):
+            stats = Location.cleanup_outside_hierarchy(db_session, dry_run=True)
+
+        # Dry run reports what would be removed but doesn't call delete_documents
+        assert stats["entities_removed"] == 1
+        mock_search_service.delete_documents.assert_not_called()

poliloom/tests/test_entity_importer.py

@@ -427,6 +427,139 @@ def test_insert_wikipedia_projects_batch_with_duplicates_handling(
         assert final_projects[0].name == "English Wikipedia Updated"
 
 
+class TestSearchServiceIndexing:
+    """Test that entity importing calls search service for label indexing.
+
+    Note: Only entities with SearchIndexedMixin are indexed (Country, Language, Location).
+    Position uses embedding-based search instead.
+    """
+
+    def test_insert_multiple_locations_indexes_labels(
+        self, db_session, mock_search_service
+    ):
+        """Test that inserting multiple locations calls index_documents with correct data."""
+        locations = [
+            {
+                "wikidata_id": "Q60",
+                "name": "New York City",
+                "description": "City in USA",
+                "labels": ["New York City", "NYC", "New York"],
+            },
+            {
+                "wikidata_id": "Q84",
+                "name": "London",
+                "description": "Capital of UK",
+                "labels": ["London", "Greater London"],
+            },
+        ]
+
+        collection = EntityCollection(model_class=Location, shared_classes=frozenset())
+        for loc in locations:
+            collection.add_entity(loc)
+
+        collection.insert(db_session, mock_search_service)
+
+        # Verify index_documents was called
+        mock_search_service.index_documents.assert_called_once()
+
+        # Verify the documents passed to index_documents
+        call_args = mock_search_service.index_documents.call_args[0][0]
+        assert len(call_args) == 2
+
+        # Check first document
+        doc1 = next(d for d in call_args if d["id"] == "Q60")
+        assert doc1["type"] == "locations"
+        assert set(doc1["labels"]) == {"New York City", "NYC", "New York"}
+
+        # Check second document
+        doc2 = next(d for d in call_args if d["id"] == "Q84")
+        assert doc2["type"] == "locations"
+        assert set(doc2["labels"]) == {"London", "Greater London"}
+
+    def test_insert_locations_indexes_labels(self, db_session, mock_search_service):
+        """Test that inserting locations calls index_documents with correct data."""
+        locations = [
+            {
+                "wikidata_id": "Q60",
+                "name": "New York City",
+                "description": "City in USA",
+                "labels": ["New York City", "NYC", "New York"],
+            },
+        ]
+
+        collection = EntityCollection(model_class=Location, shared_classes=frozenset())
+        for loc in locations:
+            collection.add_entity(loc)
+
+        collection.insert(db_session, mock_search_service)
+
+        # Verify index_documents was called with correct data
+        mock_search_service.index_documents.assert_called_once()
+        call_args = mock_search_service.index_documents.call_args[0][0]
+        assert len(call_args) == 1
+        assert call_args[0]["id"] == "Q60"
+        assert call_args[0]["type"] == "locations"
+        assert set(call_args[0]["labels"]) == {"New York City", "NYC", "New York"}
+
+    def test_insert_countries_indexes_labels(self, db_session, mock_search_service):
+        """Test that inserting countries calls index_documents with correct data."""
+        countries = [
+            {
+                "wikidata_id": "Q30",
+                "name": "United States",
+                "description": "Country",
+                "iso_code": "US",
+                "labels": ["United States", "USA", "America"],
+            },
+        ]
+
+        collection = EntityCollection(model_class=Country, shared_classes=frozenset())
+        for country in countries:
+            collection.add_entity(country)
+
+        collection.insert(db_session, mock_search_service)
+
+        # Verify index_documents was called
+        mock_search_service.index_documents.assert_called_once()
+        call_args = mock_search_service.index_documents.call_args[0][0]
+        assert call_args[0]["id"] == "Q30"
+        assert call_args[0]["type"] == "countries"
+
+    def test_insert_empty_batch_does_not_call_index(
+        self, db_session, mock_search_service
+    ):
+        """Test that empty batch does not call index_documents."""
+        collection = EntityCollection(model_class=Location, shared_classes=frozenset())
+        collection.insert(db_session, mock_search_service)
+
+        # index_documents should not be called for empty batch
+        mock_search_service.index_documents.assert_not_called()
+
+    def test_insert_entities_without_labels_indexes_empty_labels(
+        self, db_session, mock_search_service
+    ):
+        """Test that entities without labels are indexed with empty labels list."""
+        locations = [
+            {
+                "wikidata_id": "Q1",
+                "name": "Some Location",
+                "description": "No labels",
+                # No labels key
+            },
+        ]
+
+        collection = EntityCollection(model_class=Location, shared_classes=frozenset())
+        for loc in locations:
+            collection.add_entity(loc)
+
+        collection.insert(db_session, mock_search_service)
+
+        # Should still call index_documents
+        mock_search_service.index_documents.assert_called_once()
+        call_args = mock_search_service.index_documents.call_args[0][0]
+        assert call_args[0]["labels"] == []
+
+
 class TestWikipediaProjectFiltering:
     """Test Wikipedia project filtering logic in should_import method."""