tests/test_sage_v2.py

import unittest
from unittest.mock import MagicMock, patch, mock_open
import sys
import os
import json
import torch

# Add project root to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

# Mock Transformers before importing modules that use them
with patch('transformers.AutoProcessor.from_pretrained'), \
     patch('transformers.AutoModelForCausalLM.from_pretrained'), \
     patch('transformers.AutoTokenizer.from_pretrained'):
    import translation_module
    import llm_module
    import agent_module
    import oracle_module
    import ui_module

class TestSageComprehensive(unittest.TestCase):

    # --- Translation Module ---
    @patch('translation_module.get_translation')
    def test_translation(self, mock_trans):
        # Return input text to verify "Sage 6.5" is present in the final dict
        mock_trans.side_effect = lambda text, lang, cache: text
        res = translation_module.localize_init("de")
        self.assertEqual(res["lang"], "German")
        self.assertTrue("Sage 6.5" in res["welcome"])

    def test_localize_init_direct(self):
        # Without translation
        res = translation_module.localize_init("en")
        self.assertEqual(res["lang"], "English")
        self.assertTrue("Sage 6.5" in res["welcome"])
        self.assertTrue("specific Topic" in res["welcome"])

    # --- LLM Module ---
    @patch('llm_module.get_llm')
    def test_detect_language(self, mock_get_llm):
        mock_model = MagicMock()
        mock_processor = MagicMock()
        mock_get_llm.return_value = (mock_model, mock_processor)
        mock_model.device = "cpu"
        
        # Test conversational output
        mock_processor.batch_decode.return_value = ["The language used is German, I believe."]
        lang = llm_module.detect_language("Hallo wie gehts")
        self.assertEqual(lang, "German")
        
        # Test fallback
        mock_processor.batch_decode.return_value = ["Unknown."]
        lang = llm_module.detect_language("xyz")
        self.assertEqual(lang, "English")

    # --- Agent Module ---
    def test_compress_history(self):
        history = [{"role": "user", "content": "hi"}] * 20
        compressed = agent_module.compress_history(history, max_turns=5)
        self.assertEqual(len(compressed), 10)

    @patch('agent_module.get_llm')
    @patch('agent_module.TextIteratorStreamer')
    def test_agentic_tool_call(self, mock_streamer, mock_llm):
        mock_m = MagicMock()
        mock_p = MagicMock()
        mock_llm.return_value = (mock_m, mock_p)
        mock_m.device = "cpu"
        
        # Simulate LLM deciding to call tool
        # The agent logic: 
        # 1. User: "My name is Julian"
        # 2. LLM via generation: "Greetings... <tool_call>...</tool_call>"
        # 3. Agent parses this, calls oracle, appends result.
        
        # We mock the streamer to yield the tool call
        tool_json = json.dumps({"name": "oracle_consultation", "arguments": {"topic": "General", "name": "Julian", "date_str": "today"}})
        mock_stream_iter = iter([f"Thinking... <tool_call>{tool_json}</tool_call>"])
        mock_streamer.return_value = mock_stream_iter
        
        gen = agent_module.chat_agent_stream("My name is Julian", [])
        
        # Consuming the generator
        results = list(gen)
        
        # We expect:
        # 1. "*(Consulting the Oracle...)*"
        # 2. "__TURN_END__" (which signals the UI to refresh/append)
        
        self.assertTrue("*(Consulting the Oracle...)*" in results)
        self.assertTrue("__TURN_END__" in results)

    @patch('agent_module.get_llm')
    @patch('agent_module.TextIteratorStreamer')
    def test_role_alternation_fix(self, mock_streamer, mock_llm):
        mock_m = MagicMock()
        mock_p = MagicMock()
        mock_llm.return_value = (mock_m, mock_p)
        mock_m.device = "cpu"
        
        # Start with assistant message
        history = [{"role": "assistant", "content": "Welcome"}]
        
        # We want to check if apply_chat_template is called with a messages list 
        # that alternates user/assistant correctly.
        gen = agent_module.chat_agent_stream("hi", history)
        # We need to trigger a turn
        mock_it = MagicMock()
        mock_it.__iter__.return_value = ["Hello"]
        mock_streamer.return_value = mock_it
        
        list(gen)
        
        # Check call arguments
        # messages should have: [user (intro+greetings), assistant (welcome), user (hi)]
        args, kwargs = mock_p.apply_chat_template.call_args
        messages = args[0]
        self.assertEqual(messages[0]["role"], "user") # Intro
        self.assertEqual(messages[1]["role"], "assistant") # Welcome
        self.assertEqual(messages[2]["role"], "user") # Hi
        self.assertEqual(len(messages), 3)

    # --- UI Module ---
    def test_save_and_clear(self):
        empty, msg = ui_module.save_and_clear("Hello")
        self.assertEqual(empty, "")
        self.assertEqual(msg, "Hello")

    def test_clear_messages(self):
        with patch('ui_module.localize_init') as mock_loc:
            mock_loc.return_value = {"welcome": "Willkommen", "lang": "German"}
            welcome, hist = ui_module.clear_messages("German")
            self.assertEqual(welcome[0]["content"], "Willkommen")
            self.assertEqual(hist[0]["content"], "Willkommen")

    @patch('builtins.open', new_callable=mock_open, read_data='[{"role": "user", "content": "hi"}]')
    def test_import_chat(self, m):
        mock_file = MagicMock()
        mock_file.name = "test.json"
        chatbot, hist = ui_module.import_chat(mock_file)
        self.assertEqual(len(chatbot), 1)
        self.assertEqual(chatbot[0]["content"], "hi")

    def test_ui_wiring(self):
        demo = ui_module.build_demo()
        self.assertTrue(len(demo.fns) > 5)

    @patch('ui_module.chat_agent_stream')
    def test_chat_wrapper(self, mock_stream):
        # Part1, Part2, Part3
        mock_stream.return_value = iter(["Part 1", "Part 2", "__TURN_END__", "Final"])
        history = []
        gen = ui_module.chat_wrapper("Hello", history)
        results = list(gen)
        final_chatbot, final_hist = results[-1]
        # [User, Assistant1(Part2), Assistant2(Final)]
        self.assertEqual(len(final_chatbot), 3)
        self.assertEqual(final_chatbot[2]["role"], "assistant")
        self.assertEqual(final_chatbot[2]["content"], "Final")

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