import unittest
from unittest.mock import MagicMock
from rlm.cognitive_tools import CognitiveTools  # Assuming cognitive_tools.py exists in rlm package

class TestCognitiveTools(unittest.TestCase):

    def setUp(self):
        self.llm = MagicMock()  # Mock LLM for testing
        self.cognitive_tools = CognitiveTools(self.llm)

    def test_generate_code(self):
        self.llm.predict.return_value = "def add(x, y): return x + y"
        code = self.cognitive_tools.generate_code("python", "add two numbers")
        self.assertEqual(code, "def add(x, y): return x + y")
        self.llm.predict.assert_called_once()

    def test_generate_tests(self):
        self.llm.predict.return_value = "def test_add(): assert add(1, 2) == 3"
        tests = self.cognitive_tools.generate_tests("python", "add(x, y)")
        self.assertEqual(tests, "def test_add(): assert add(1, 2) == 3")
        self.llm.predict.assert_called_once()

    def test_analyze_code(self):
        self.llm.predict.return_value = "This code adds two numbers."
        analysis = self.cognitive_tools.analyze_code("python", "def add(x, y): return x + y")
        self.assertEqual(analysis, "This code adds two numbers.")
        self.llm.predict.assert_called_once()

    def test_improve_code(self):
        self.llm.predict.return_value = "def add(x, y):\n  '''Adds two numbers.'''\n  return x + y"
        improved_code = self.cognitive_tools.improve_code("python", "def add(x, y): return x + y")
        self.assertEqual(improved_code, "def add(x, y):\n  '''Adds two numbers.'''\n  return x + y")
        self.llm.predict.assert_called_once()

if __name__ == '__main__':
    unittest.main()