33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
import os
|
|
import unittest
|
|
|
|
from src.agent import RAGAgent
|
|
|
|
class TestRAGAgent(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
# Ensure data directory exists with at least one document
|
|
data_dir = "data"
|
|
os.makedirs(data_dir, exist_ok=True)
|
|
sample_path = os.path.join(data_dir, "sample.txt")
|
|
with open(sample_path, "w", encoding="utf-8") as f:
|
|
f.write("Python is a versatile programming language used for web development, data science, and automation.")
|
|
cls.agent = RAGAgent(config_path="src/config.yaml")
|
|
|
|
def test_retrieve_non_empty(self):
|
|
passages = self.agent.kb.retrieve("Python programming", top_k=2)
|
|
self.assertTrue(len(passages) > 0)
|
|
self.assertIn("Python is a versatile programming language", passages[0][0])
|
|
|
|
def test_generate_response(self):
|
|
answer = self.agent.generate_response("What is Python?")
|
|
self.assertIsInstance(answer, str)
|
|
self.assertTrue(len(answer) > 0)
|
|
|
|
def test_empty_query(self):
|
|
answer = self.agent.generate_response("")
|
|
self.assertIsInstance(answer, str)
|
|
self.assertIn("No relevant information found", answer)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |