| 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() |
|
|
