import time
import json
from datetime import datetime
from loguru import logger
from swarms.communication.redis_wrap import (
RedisConversation,
REDIS_AVAILABLE,
)
class TestResults:
def __init__(self):
self.results = []
self.start_time = datetime.now()
self.end_time = None
self.total_tests = 0
self.passed_tests = 0
self.failed_tests = 0
def add_result(
self, test_name: str, passed: bool, error: str = None
):
self.total_tests += 1
if passed:
self.passed_tests += 1
status = "✅ PASSED"
else:
self.failed_tests += 1
status = "❌ FAILED"
self.results.append(
{
"test_name": test_name,
"status": status,
"error": error if error else "None",
}
)
def generate_markdown(self) -> str:
self.end_time = datetime.now()
duration = (self.end_time - self.start_time).total_seconds()
md = [
"# Redis Conversation Test Results",
"",
f"Test Run: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}",
f"Duration: {duration:.2f} seconds",
"",
"## Summary",
f"- Total Tests: {self.total_tests}",
f"- Passed: {self.passed_tests}",
f"- Failed: {self.failed_tests}",
f"- Success Rate: {(self.passed_tests/self.total_tests*100):.1f}%",
"",
"## Detailed Results",
"",
"| Test Name | Status | Error |",
"|-----------|--------|-------|",
]
for result in self.results:
md.append(
f"| {result['test_name']} | {result['status']} | {result['error']} |"
)
return "\n".join(md)
class RedisConversationTester:
def __init__(self):
self.results = TestResults()
self.conversation = None
self.redis_server = None
def run_test(self, test_func: callable, test_name: str):
"""Run a single test and record its result."""
try:
test_func()
self.results.add_result(test_name, True)
except Exception as e:
self.results.add_result(test_name, False, str(e))
logger.error(f"Test '{test_name}' failed: {str(e)}")
def setup(self):
"""Initialize Redis server and conversation for testing."""
try:
# # Start embedded Redis server
# self.redis_server = EmbeddedRedis(port=6379)
# if not self.redis_server.start():
# logger.error("Failed to start embedded Redis server")
# return False
# Initialize Redis conversation
self.conversation = RedisConversation(
system_prompt="Test System Prompt",
redis_host="localhost",
redis_port=6379,
redis_retry_attempts=3,
use_embedded_redis=True,
)
return True
except Exception as e:
logger.error(
f"Failed to initialize Redis conversation: {str(e)}"
)
return False
def cleanup(self):
"""Cleanup resources after tests."""
if self.redis_server:
self.redis_server.stop()
def test_initialization(self):
"""Test basic initialization."""
assert (
self.conversation is not None
), "Failed to initialize RedisConversation"
assert (
self.conversation.system_prompt == "Test System Prompt"
), "System prompt not set correctly"
def test_add_message(self):
"""Test adding messages."""
self.conversation.add("user", "Hello")
self.conversation.add("assistant", "Hi there!")
messages = self.conversation.return_messages_as_list()
assert len(messages) >= 2, "Failed to add messages"
def test_json_message(self):
"""Test adding JSON messages."""
json_content = {"key": "value", "nested": {"data": 123}}
self.conversation.add("system", json_content)
last_message = self.conversation.get_final_message_content()
assert isinstance(
json.loads(last_message), dict
), "Failed to handle JSON message"
def test_search(self):
"""Test search functionality."""
self.conversation.add("user", "searchable message")
results = self.conversation.search("searchable")
assert len(results) > 0, "Search failed to find message"
def test_delete(self):
"""Test message deletion."""
initial_count = len(
self.conversation.return_messages_as_list()
)
self.conversation.delete(0)
new_count = len(self.conversation.return_messages_as_list())
assert (
new_count == initial_count - 1
), "Failed to delete message"
def test_update(self):
"""Test message update."""
# Add initial message
self.conversation.add("user", "original message")
# Update the message
self.conversation.update(0, "user", "updated message")
# Get the message directly using query
updated_message = self.conversation.query(0)
# Verify the update
assert (
updated_message["content"] == "updated message"
), "Message content should be updated"
def test_clear(self):
"""Test clearing conversation."""
self.conversation.add("user", "test message")
self.conversation.clear()
messages = self.conversation.return_messages_as_list()
assert len(messages) == 0, "Failed to clear conversation"
def test_export_import(self):
"""Test export and import functionality."""
self.conversation.add("user", "export test")
self.conversation.export_conversation("test_export.txt")
self.conversation.clear()
self.conversation.import_conversation("test_export.txt")
messages = self.conversation.return_messages_as_list()
assert (
len(messages) > 0
), "Failed to export/import conversation"
def test_json_operations(self):
"""Test JSON operations."""
self.conversation.add("user", "json test")
json_data = self.conversation.to_json()
assert isinstance(
json.loads(json_data), list
), "Failed to convert to JSON"
def test_yaml_operations(self):
"""Test YAML operations."""
self.conversation.add("user", "yaml test")
yaml_data = self.conversation.to_yaml()
assert isinstance(yaml_data, str), "Failed to convert to YAML"
def test_token_counting(self):
"""Test token counting functionality."""
self.conversation.add("user", "token test message")
time.sleep(1) # Wait for async token counting
messages = self.conversation.to_dict()
assert any(
"token_count" in msg for msg in messages
), "Failed to count tokens"
def test_cache_operations(self):
"""Test cache operations."""
self.conversation.add("user", "cache test")
stats = self.conversation.get_cache_stats()
assert isinstance(stats, dict), "Failed to get cache stats"
def test_conversation_stats(self):
"""Test conversation statistics."""
self.conversation.add("user", "stats test")
counts = self.conversation.count_messages_by_role()
assert isinstance(
counts, dict
), "Failed to get message counts"
def run_all_tests(self):
"""Run all tests and generate report."""
if not REDIS_AVAILABLE:
logger.error(
"Redis is not available. Please install redis package."
)
return "# Redis Tests Failed\n\nRedis package is not installed."
try:
if not self.setup():
logger.error("Failed to setup Redis connection.")
return "# Redis Tests Failed\n\nFailed to connect to Redis server."
tests = [
(self.test_initialization, "Initialization Test"),
(self.test_add_message, "Add Message Test"),
(self.test_json_message, "JSON Message Test"),
(self.test_search, "Search Test"),
(self.test_delete, "Delete Test"),
(self.test_update, "Update Test"),
(self.test_clear, "Clear Test"),
(self.test_export_import, "Export/Import Test"),
(self.test_json_operations, "JSON Operations Test"),
(self.test_yaml_operations, "YAML Operations Test"),
(self.test_token_counting, "Token Counting Test"),
(self.test_cache_operations, "Cache Operations Test"),
(
self.test_conversation_stats,
"Conversation Stats Test",
),
]
for test_func, test_name in tests:
self.run_test(test_func, test_name)
return self.results.generate_markdown()
finally:
self.cleanup()
def main():
"""Main function to run tests and save results."""
tester = RedisConversationTester()
markdown_results = tester.run_all_tests()
# Save results to file
with open("redis_test_results.md", "w") as f:
f.write(markdown_results)
logger.info(
"Test results have been saved to redis_test_results.md"
)
if __name__ == "__main__":
main()