diff options
Diffstat (limited to 'autogpts/forge/forge/memory/memstore_test.py')
-rw-r--r-- | autogpts/forge/forge/memory/memstore_test.py | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/autogpts/forge/forge/memory/memstore_test.py b/autogpts/forge/forge/memory/memstore_test.py new file mode 100644 index 000000000..3eb24309b --- /dev/null +++ b/autogpts/forge/forge/memory/memstore_test.py @@ -0,0 +1,58 @@ +import hashlib +import shutil + +import pytest + +from forge.sdk.memory.memstore import ChromaMemStore + + +@pytest.fixture +def memstore(): + mem = ChromaMemStore(".test_mem_store") + yield mem + shutil.rmtree(".test_mem_store") + + +def test_add(memstore): + task_id = "test_task" + document = "This is a test document." + metadatas = {"metadata": "test_metadata"} + memstore.add(task_id, document, metadatas) + doc_id = hashlib.sha256(document.encode()).hexdigest()[:20] + assert memstore.client.get_or_create_collection(task_id).count() == 1 + + +def test_query(memstore): + task_id = "test_task" + document = "This is a test document." + metadatas = {"metadata": "test_metadata"} + memstore.add(task_id, document, metadatas) + query = "test" + assert len(memstore.query(task_id, query)["documents"]) == 1 + + +def test_update(memstore): + task_id = "test_task" + document = "This is a test document." + metadatas = {"metadata": "test_metadata"} + memstore.add(task_id, document, metadatas) + doc_id = hashlib.sha256(document.encode()).hexdigest()[:20] + updated_document = "This is an updated test document." + updated_metadatas = {"metadata": "updated_test_metadata"} + memstore.update(task_id, [doc_id], [updated_document], [updated_metadatas]) + assert memstore.get(task_id, [doc_id]) == { + "documents": [updated_document], + "metadatas": [updated_metadatas], + "embeddings": None, + "ids": [doc_id], + } + + +def test_delete(memstore): + task_id = "test_task" + document = "This is a test document." + metadatas = {"metadata": "test_metadata"} + memstore.add(task_id, document, metadatas) + doc_id = hashlib.sha256(document.encode()).hexdigest()[:20] + memstore.delete(task_id, doc_id) + assert memstore.client.get_or_create_collection(task_id).count() == 0 |