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