from __future__ import annotations

import os
import unittest

from src.model_status import ModelStatus, statuses_with_llm_status, warm_llm_model


class ModelStatusTests(unittest.TestCase):
    def test_warm_llm_model_skips_remote_when_mock_enabled(self) -> None:
        old_value = os.environ.get("VC_USE_MOCK")
        os.environ["VC_USE_MOCK"] = "1"
        try:
            status = warm_llm_model(timeout_s=0.01)
        finally:
            if old_value is None:
                os.environ.pop("VC_USE_MOCK", None)
            else:
                os.environ["VC_USE_MOCK"] = old_value

        self.assertEqual(status.kind, "llm")
        self.assertEqual(status.state, "mock")
        self.assertIn("mock", status.message.lower())

    def test_statuses_with_llm_status_replaces_initial_llm_row(self) -> None:
        llm_status = ModelStatus("llm", "loading", "载入中", message="正在启动主模型")

        statuses = statuses_with_llm_status(llm_status)

        self.assertGreaterEqual(len(statuses), 1)
        self.assertEqual(statuses[0], llm_status)
        self.assertEqual(statuses[0].message, "正在启动主模型")


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