Fix / cleanup ai indexing test

This commit is contained in:
shamoon 2025-04-29 23:51:47 -07:00
parent c74c483674
commit 78553ca0d6
No known key found for this signature in database

View File

@ -138,36 +138,71 @@ def test_load_or_build_index_builds_when_nodes_given(
temp_llm_index_dir, temp_llm_index_dir,
real_document, real_document,
): ):
with patch( with (
"paperless_ai.indexing.load_index_from_storage", patch(
side_effect=ValueError("Index not found"), "paperless_ai.indexing.load_index_from_storage",
): side_effect=ValueError("Index not found"),
with patch( ),
patch(
"paperless_ai.indexing.VectorStoreIndex", "paperless_ai.indexing.VectorStoreIndex",
return_value=MagicMock(), return_value=MagicMock(),
) as mock_index_cls: ) as mock_index_cls,
with patch( patch(
"paperless_ai.indexing.get_or_create_storage_context", "paperless_ai.indexing.get_or_create_storage_context",
return_value=MagicMock(), return_value=MagicMock(),
) as mock_storage: ) as mock_storage,
mock_storage.return_value.persist_dir = temp_llm_index_dir ):
indexing.load_or_build_index( mock_storage.return_value.persist_dir = temp_llm_index_dir
nodes=[indexing.build_document_node(real_document)], indexing.load_or_build_index(
) nodes=[indexing.build_document_node(real_document)],
mock_index_cls.assert_called_once() )
mock_index_cls.assert_called_once()
def test_load_or_build_index_raises_exception_when_no_nodes( def test_load_or_build_index_raises_exception_when_no_nodes(
temp_llm_index_dir, temp_llm_index_dir,
mock_embed_model,
): ):
with patch( with (
"paperless_ai.indexing.load_index_from_storage", patch(
side_effect=ValueError("Index not found"), "paperless_ai.indexing.load_index_from_storage",
side_effect=ValueError("Index not found"),
),
patch(
"paperless_ai.indexing.get_or_create_storage_context",
return_value=MagicMock(),
),
): ):
with pytest.raises(Exception): with pytest.raises(Exception):
indexing.load_or_build_index() indexing.load_or_build_index()
@pytest.mark.django_db
def test_load_or_build_index_succeeds_when_nodes_given(
temp_llm_index_dir,
mock_embed_model,
):
with (
patch(
"paperless_ai.indexing.load_index_from_storage",
side_effect=ValueError("Index not found"),
),
patch(
"paperless_ai.indexing.VectorStoreIndex",
return_value=MagicMock(),
) as mock_index_cls,
patch(
"paperless_ai.indexing.get_or_create_storage_context",
return_value=MagicMock(),
) as mock_storage,
):
mock_storage.return_value.persist_dir = temp_llm_index_dir
indexing.load_or_build_index(
nodes=[MagicMock()],
)
mock_index_cls.assert_called_once()
@pytest.mark.django_db @pytest.mark.django_db
def test_add_or_update_document_updates_existing_entry( def test_add_or_update_document_updates_existing_entry(
temp_llm_index_dir, temp_llm_index_dir,