|
|
|
"""Tests for the MilvusMemory class.""" |
|
import os |
|
import sys |
|
import unittest |
|
|
|
try: |
|
from autogpt.memory.milvus import MilvusMemory |
|
|
|
def mock_config() -> dict: |
|
"""Mock the Config class""" |
|
return type( |
|
"MockConfig", |
|
(object,), |
|
{ |
|
"debug_mode": False, |
|
"continuous_mode": False, |
|
"speak_mode": False, |
|
"milvus_collection": "autogpt", |
|
"milvus_addr": "localhost:19530", |
|
}, |
|
) |
|
|
|
class TestMilvusMemory(unittest.TestCase): |
|
"""Tests for the MilvusMemory class.""" |
|
|
|
def setUp(self) -> None: |
|
"""Set up the test environment""" |
|
self.cfg = mock_config() |
|
self.memory = MilvusMemory(self.cfg) |
|
|
|
def test_add(self) -> None: |
|
"""Test adding a text to the cache""" |
|
text = "Sample text" |
|
self.memory.clear() |
|
self.memory.add(text) |
|
result = self.memory.get(text) |
|
self.assertEqual([text], result) |
|
|
|
def test_clear(self) -> None: |
|
"""Test clearing the cache""" |
|
self.memory.clear() |
|
self.assertEqual(self.memory.collection.num_entities, 0) |
|
|
|
def test_get(self) -> None: |
|
"""Test getting a text from the cache""" |
|
text = "Sample text" |
|
self.memory.clear() |
|
self.memory.add(text) |
|
result = self.memory.get(text) |
|
self.assertEqual(result, [text]) |
|
|
|
def test_get_relevant(self) -> None: |
|
"""Test getting relevant texts from the cache""" |
|
text1 = "Sample text 1" |
|
text2 = "Sample text 2" |
|
self.memory.clear() |
|
self.memory.add(text1) |
|
self.memory.add(text2) |
|
result = self.memory.get_relevant(text1, 1) |
|
self.assertEqual(result, [text1]) |
|
|
|
def test_get_stats(self) -> None: |
|
"""Test getting the cache stats""" |
|
text = "Sample text" |
|
self.memory.clear() |
|
self.memory.add(text) |
|
stats = self.memory.get_stats() |
|
self.assertEqual(15, len(stats)) |
|
|
|
except: |
|
print("Milvus not installed, skipping tests") |
|
|
