cognee/level_2/tests/crud_test.py

78 lines
3.3 KiB
Python

import unittest
import asyncio
import sys
sys.path.append("..") # Adds higher directory to python modules path.
from level_2.level_2_pdf_vectorstore__dlt_contracts import Memory
class TestMemory(unittest.TestCase):
def setUp(self):
self.loop = asyncio.get_event_loop()
self.memory = Memory(user_id="123")
self.loop.run_until_complete(self.memory.async_init())
def test_add_fetch_delete_semantic_memory(self):
async def semantic_workflow():
params = {"sample_param": "value"}
sample_memory = "sample semantic memory"
# Add
await self.memory._add_semantic_memory(sample_memory, params=params)
# Fetch
fetched = await self.memory._fetch_semantic_memory(sample_memory, params)
fetched_text = fetched['data']['Get']['EPISODICMEMORY'][0]['text']
self.assertIn(sample_memory, fetched_text) # Replace this with the appropriate validation
# Delete
await self.memory._delete_semantic_memory()
# Verify Deletion
after_delete = await self.memory._fetch_semantic_memory(sample_memory, params)
self.assertNotIn(sample_memory, after_delete) # Replace with the appropriate validation
self.loop.run_until_complete(semantic_workflow())
def test_add_fetch_delete_episodic_memory(self):
async def episodic_workflow():
params = {"sample_param": "value"}
sample_memory = """{
"sample_key": "sample_value"
}"""
# Add
await self.memory._add_episodic_memory(observation=sample_memory, params=params)
# Fetch
fetched = await self.memory._fetch_episodic_memory(sample_memory)
fetched_text = fetched['data']['Get']['EPISODICMEMORY'][0]['text']
self.assertIn(sample_memory, fetched_text) # Replace this with the appropriate validation
# Delete
await self.memory._delete_episodic_memory()
# Verify Deletion
after_delete = await self.memory._fetch_episodic_memory(sample_memory)
self.assertNotIn(sample_memory, after_delete) # Replace with the appropriate validation
self.loop.run_until_complete(episodic_workflow())
# def test_add_fetch_delete_buffer_memory(self):
# async def buffer_workflow():
# params = {"sample_param": "value"}
# user_input = "sample buffer input"
# namespace = "sample_namespace"
#
# # Add
# await self.memory._add_buffer_memory(user_input=user_input, namespace=namespace, params=params)
# # Fetch
# fetched = await self.memory._fetch_buffer_memory(user_input, namespace)
# self.assertIn(user_input, fetched) # Replace this with the appropriate validation
# # Delete
# await self.memory._delete_buffer_memory()
# # Verify Deletion
# after_delete = await self.memory._fetch_buffer_memory(user_input, namespace)
# self.assertNotIn(user_input, after_delete) # Replace with the appropriate validation
#
# self.loop.run_until_complete(buffer_workflow())
if __name__ == '__main__':
unittest.main()