import unittest

from itertools import tee

from . import test_data
from . import test_util

from sia.llm_engine import LlmEngine

class LlmEngineTest(unittest.TestCase):
    def setUp(self):
        self.model_path = "/root/model"

    def test_initialization(self):
        llm_engine = LlmEngine(self.model_path)
        self.assertIsInstance(llm_engine, LlmEngine)

    def test_infer(self):
        main_context = "This is a test"
        llm_engine = LlmEngine(self.model_path)
        tokens = llm_engine.infer(test_data.echo_system_prompt, main_context)
        print_tokens, result_tokens = tee(tokens)
        for token in print_tokens:
            print(token, end="", flush=True)
        result = ''.join(result_tokens)
        self.assertEqual(result, f"{main_context}<test_tag>{main_context}</test_tag>")