28 lines
940 B
Python
28 lines
940 B
Python
from unittest.mock import PropertyMock, patch
|
|
|
|
from llama_index import Document
|
|
|
|
from private_gpt.server.ingest.ingest_service import IngestService
|
|
from tests.fixtures.mock_injector import MockInjector
|
|
|
|
|
|
def test_save_many_nodes(injector: MockInjector) -> None:
|
|
"""This is a specific test for a local Chromadb Vector Database setup.
|
|
|
|
Extend it when we add support for other vector databases in VectorStoreComponent.
|
|
"""
|
|
with patch(
|
|
"chromadb.api.segment.SegmentAPI.max_batch_size", new_callable=PropertyMock
|
|
) as max_batch_size:
|
|
# Make max batch size of Chromadb very small
|
|
max_batch_size.return_value = 10
|
|
|
|
ingest_service = injector.get(IngestService)
|
|
|
|
documents = []
|
|
for _i in range(100):
|
|
documents.append(Document(text="This is a sentence."))
|
|
|
|
ingested_docs = ingest_service._save_docs(documents)
|
|
assert len(ingested_docs) == len(documents)
|