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()