315 lines
8.6 KiB
Python
315 lines
8.6 KiB
Python
"""
|
|
Tests for IngestionService - document ingestion operations.
|
|
|
|
Tests cover:
|
|
- Single page ingestion
|
|
- Batch ingestion
|
|
- Full re-index (ingest_all_pages)
|
|
- list_all_pages usage
|
|
|
|
Run with: pytest tests/test_ingestion.py -v -s
|
|
"""
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from typing import AsyncGenerator
|
|
|
|
from src.services.ingestion_service import IngestionService
|
|
from src.models.ingestion import (
|
|
IngestionRequest,
|
|
IngestionResult,
|
|
BatchIngestionRequest,
|
|
BatchIngestionResult
|
|
)
|
|
|
|
|
|
# Test constants
|
|
TEST_USER = "ingestion-tester"
|
|
TEST_PAGE_ID = 456
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_vector_service():
|
|
"""Mock Vector service."""
|
|
mock = AsyncMock()
|
|
mock.update_from_page = AsyncMock(return_value=MagicMock(
|
|
chunks_created=2,
|
|
chunks_deleted=0,
|
|
success=True
|
|
))
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_graph_service():
|
|
"""Mock Graph service."""
|
|
mock = AsyncMock()
|
|
mock.update_from_page = AsyncMock(return_value=MagicMock(
|
|
entities_extracted=5,
|
|
relationships_created=5,
|
|
success=True
|
|
))
|
|
mock.create_entity_mention_links = AsyncMock(return_value=3)
|
|
mock.get_all_entities = AsyncMock(return_value=[])
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_wiki_client():
|
|
"""Mock Wiki.js client."""
|
|
mock = AsyncMock()
|
|
mock.get_page = AsyncMock(return_value={
|
|
"id": TEST_PAGE_ID,
|
|
"title": "Test Page",
|
|
"path": f"users/{TEST_USER}/test",
|
|
"content": "Test content here.",
|
|
"tags": ["test"]
|
|
})
|
|
mock.list_all_pages = AsyncMock(return_value=[
|
|
{"id": 1, "path": f"users/{TEST_USER}/page1", "title": "Page 1"},
|
|
{"id": 2, "path": f"users/{TEST_USER}/page2", "title": "Page 2"},
|
|
{"id": 3, "path": f"users/{TEST_USER}/page3", "title": "Page 3"},
|
|
])
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def ingestion_service(mock_vector_service, mock_graph_service, mock_wiki_client):
|
|
"""Get IngestionService with mocked dependencies."""
|
|
return IngestionService(
|
|
vector_service=mock_vector_service,
|
|
graph_service=mock_graph_service,
|
|
wiki_client=mock_wiki_client
|
|
)
|
|
|
|
|
|
class TestSinglePageIngestion:
|
|
"""Test single page ingestion."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_page_success(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test successful page ingestion."""
|
|
result = await ingestion_service.ingest_page(
|
|
page_id=TEST_PAGE_ID,
|
|
user=TEST_USER
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.page_id == TEST_PAGE_ID
|
|
mock_wiki_client.get_page.assert_called_once_with(TEST_PAGE_ID)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_page_not_found(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test ingestion when page not found."""
|
|
mock_wiki_client.get_page.return_value = None
|
|
|
|
result = await ingestion_service.ingest_page(
|
|
page_id=999,
|
|
user=TEST_USER
|
|
)
|
|
|
|
assert result.success is False
|
|
assert "not found" in result.error.lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_page_skip_vectors(
|
|
self,
|
|
ingestion_service,
|
|
mock_vector_service,
|
|
mock_graph_service
|
|
):
|
|
"""Test ingestion with vectors skipped."""
|
|
result = await ingestion_service.ingest_page(
|
|
page_id=TEST_PAGE_ID,
|
|
user=TEST_USER,
|
|
skip_vectors=True
|
|
)
|
|
|
|
assert result.success is True
|
|
# Vector service should not be called
|
|
mock_vector_service.update_from_page.assert_not_called()
|
|
# Graph service should still be called
|
|
mock_graph_service.update_from_page.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_page_skip_graph(
|
|
self,
|
|
ingestion_service,
|
|
mock_vector_service,
|
|
mock_graph_service
|
|
):
|
|
"""Test ingestion with graph skipped."""
|
|
result = await ingestion_service.ingest_page(
|
|
page_id=TEST_PAGE_ID,
|
|
user=TEST_USER,
|
|
skip_graph=True
|
|
)
|
|
|
|
assert result.success is True
|
|
# Vector service should be called
|
|
mock_vector_service.update_from_page.assert_called_once()
|
|
# Graph service should not be called
|
|
mock_graph_service.update_from_page.assert_not_called()
|
|
|
|
|
|
class TestBatchIngestion:
|
|
"""Test batch page ingestion."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_batch_success(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test successful batch ingestion."""
|
|
result = await ingestion_service.ingest_batch(
|
|
page_ids=[1, 2, 3],
|
|
user=TEST_USER,
|
|
max_concurrent=2
|
|
)
|
|
|
|
assert result.total_pages == 3
|
|
assert result.successful == 3
|
|
assert result.failed == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_batch_with_failures(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test batch ingestion with some failures."""
|
|
# Make page 2 not found
|
|
def get_page_side_effect(page_id):
|
|
if page_id == 2:
|
|
return None
|
|
return {
|
|
"id": page_id,
|
|
"title": f"Page {page_id}",
|
|
"path": f"users/{TEST_USER}/page{page_id}",
|
|
"content": "Content",
|
|
"tags": []
|
|
}
|
|
|
|
mock_wiki_client.get_page.side_effect = get_page_side_effect
|
|
|
|
result = await ingestion_service.ingest_batch(
|
|
page_ids=[1, 2, 3],
|
|
user=TEST_USER
|
|
)
|
|
|
|
assert result.total_pages == 3
|
|
assert result.successful == 2
|
|
assert result.failed == 1
|
|
|
|
|
|
class TestIngestAllPages:
|
|
"""Test full re-index (ingest_all_pages)."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_all_uses_list_all_pages(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test that ingest_all_pages uses list_all_pages (not search)."""
|
|
result = await ingestion_service.ingest_all_pages(
|
|
user=TEST_USER
|
|
)
|
|
|
|
# Should use list_all_pages, not search_pages
|
|
mock_wiki_client.list_all_pages.assert_called_once()
|
|
# Should have processed 3 pages from the mock
|
|
assert result.total_pages == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_all_with_path_prefix(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test ingest_all_pages with path prefix filter."""
|
|
await ingestion_service.ingest_all_pages(
|
|
user=TEST_USER,
|
|
path_prefix=f"users/{TEST_USER}/technology"
|
|
)
|
|
|
|
mock_wiki_client.list_all_pages.assert_called_once_with(
|
|
path_prefix=f"users/{TEST_USER}/technology"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_all_empty_wiki(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test ingest_all_pages when no pages found."""
|
|
mock_wiki_client.list_all_pages.return_value = []
|
|
|
|
result = await ingestion_service.ingest_all_pages(
|
|
user=TEST_USER
|
|
)
|
|
|
|
assert result.total_pages == 0
|
|
assert result.successful == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ingest_all_respects_max_concurrent(
|
|
self,
|
|
ingestion_service,
|
|
mock_wiki_client
|
|
):
|
|
"""Test that max_concurrent parameter is passed through."""
|
|
# Create many pages
|
|
mock_wiki_client.list_all_pages.return_value = [
|
|
{"id": i, "path": f"users/{TEST_USER}/page{i}", "title": f"Page {i}"}
|
|
for i in range(20)
|
|
]
|
|
|
|
result = await ingestion_service.ingest_all_pages(
|
|
user=TEST_USER,
|
|
max_concurrent=5
|
|
)
|
|
|
|
assert result.total_pages == 20
|
|
|
|
|
|
class TestIngestionModels:
|
|
"""Test ingestion request/response models."""
|
|
|
|
def test_ingestion_request_defaults(self):
|
|
"""Test IngestionRequest default values."""
|
|
request = IngestionRequest(
|
|
page_id=123,
|
|
user="testuser"
|
|
)
|
|
assert request.page_id == 123
|
|
assert request.user == "testuser"
|
|
assert request.force_refresh is False
|
|
assert request.skip_vectors is False
|
|
assert request.skip_graph is False
|
|
|
|
def test_batch_ingestion_request(self):
|
|
"""Test BatchIngestionRequest."""
|
|
request = BatchIngestionRequest(
|
|
page_ids=[1, 2, 3],
|
|
user="testuser",
|
|
max_concurrent=5
|
|
)
|
|
assert len(request.page_ids) == 3
|
|
assert request.max_concurrent == 5
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v", "-s"])
|