replace existing file records on re-import
This commit is contained in:
@@ -0,0 +1,67 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
from chromy.utilities import ingest_file
|
||||
|
||||
|
||||
class UtilityTests(unittest.TestCase):
|
||||
def test_ingest_file_adds_new_file_without_deleting(self) -> None:
|
||||
chunks = ["chunk 1", "chunk 2"]
|
||||
embeddings = [
|
||||
{"text": "chunk 1", "embedding": [0.1, 0.2]},
|
||||
{"text": "chunk 2", "embedding": [0.3, 0.4]},
|
||||
]
|
||||
|
||||
with (
|
||||
patch("chromy.utilities.has_data_for_file", return_value=False) as has_data,
|
||||
patch("chromy.utilities.delete_data") as delete_data,
|
||||
patch("chromy.utilities.chunk_file", return_value=chunks) as chunk_file,
|
||||
patch("chromy.utilities.embed", return_value=embeddings) as embed,
|
||||
patch("chromy.utilities.add_data") as add_data,
|
||||
):
|
||||
records_added = ingest_file("notes", "/tmp/play.txt")
|
||||
|
||||
has_data.assert_called_once_with("notes", "/tmp/play.txt")
|
||||
delete_data.assert_not_called()
|
||||
chunk_file.assert_called_once_with("/tmp/play.txt")
|
||||
embed.assert_called_once_with(chunks)
|
||||
add_data.assert_called_once_with("notes", embeddings, "/tmp/play.txt")
|
||||
self.assertEqual(records_added, 2)
|
||||
|
||||
def test_ingest_file_replaces_existing_file_records_before_adding(self) -> None:
|
||||
chunks = ["chunk 1"]
|
||||
embeddings = [{"text": "chunk 1", "embedding": [0.1, 0.2]}]
|
||||
manager = MagicMock()
|
||||
|
||||
with (
|
||||
patch("chromy.utilities.has_data_for_file", return_value=True) as has_data,
|
||||
patch("chromy.utilities.delete_data") as delete_data,
|
||||
patch("chromy.utilities.chunk_file", return_value=chunks) as chunk_file,
|
||||
patch("chromy.utilities.embed", return_value=embeddings) as embed,
|
||||
patch("chromy.utilities.add_data") as add_data,
|
||||
):
|
||||
manager.attach_mock(has_data, "has_data")
|
||||
manager.attach_mock(delete_data, "delete_data")
|
||||
manager.attach_mock(chunk_file, "chunk_file")
|
||||
manager.attach_mock(embed, "embed")
|
||||
manager.attach_mock(add_data, "add_data")
|
||||
|
||||
records_added = ingest_file("notes", "/tmp/play.txt")
|
||||
|
||||
self.assertEqual(
|
||||
manager.mock_calls,
|
||||
[
|
||||
call.has_data("notes", "/tmp/play.txt"),
|
||||
call.delete_data("notes", {"file_name": "/tmp/play.txt"}),
|
||||
call.chunk_file("/tmp/play.txt"),
|
||||
call.embed(chunks),
|
||||
call.add_data("notes", embeddings, "/tmp/play.txt"),
|
||||
],
|
||||
)
|
||||
self.assertEqual(records_added, 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user