Spaces:
Sleeping
Sleeping
File size: 6,509 Bytes
058ae1e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | 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()
|