| from __future__ import annotations |
|
|
| from typing import Any |
| from unittest.mock import patch |
| import subprocess |
| import unittest |
|
|
| from agent_harness.lm_studio_management import LMStudioResidencyManager, LMStudioServer |
|
|
|
|
| class FakeResidencyManager(LMStudioResidencyManager): |
| def __init__(self) -> None: |
| super().__init__("http://127.0.0.1:1234", "UNUSED") |
| self.instances: list[dict[str, Any]] = [] |
|
|
| def models(self) -> tuple[dict[str, Any], ...]: |
| by_key: dict[str, list[dict[str, Any]]] = {} |
| for item in self.instances: |
| by_key.setdefault(item["model_key"], []).append( |
| {"id": item["instance_id"], "config": item["config"]} |
| ) |
| return tuple( |
| {"key": key, "type": "llm", "loaded_instances": values} |
| for key, values in by_key.items() |
| ) |
|
|
| def _unload(self, instance_id: str) -> dict[str, Any]: |
| self.instances = [item for item in self.instances if item["instance_id"] != instance_id] |
| return {"instance_id": instance_id} |
|
|
| def _request( |
| self, method: str, endpoint: str, payload: dict[str, Any] | None = None |
| ) -> dict[str, Any]: |
| if endpoint != "/api/v1/models/load" or payload is None: |
| raise AssertionError((method, endpoint, payload)) |
| key = str(payload["model"]) |
| context = int(payload["context_length"]) |
| self.instances.append( |
| {"model_key": key, "instance_id": key, "config": {"context_length": context}} |
| ) |
| return { |
| "type": "llm", |
| "instance_id": key, |
| "status": "loaded", |
| "load_time_seconds": 1.0, |
| "load_config": {"context_length": context}, |
| } |
|
|
|
|
| class ResidencyManagerTests(unittest.TestCase): |
| def test_server_status_accepts_success_message_on_stderr(self) -> None: |
| server = LMStudioServer(cli_path=__file__) |
| completed = subprocess.CompletedProcess( |
| args=["lms", "server", "status"], |
| returncode=0, |
| stdout="", |
| stderr="The server is running on port 1234.\n", |
| ) |
| with patch.object(server, "_run", return_value=completed): |
| self.assertTrue(server.status()["running"]) |
|
|
| def test_server_status_does_not_treat_not_running_as_running(self) -> None: |
| server = LMStudioServer(cli_path=__file__) |
| completed = subprocess.CompletedProcess( |
| args=["lms", "server", "status"], |
| returncode=0, |
| stdout="", |
| stderr="The server is not running.\n", |
| ) |
| with patch.object(server, "_run", return_value=completed): |
| self.assertFalse(server.status()["running"]) |
|
|
| def test_ensure_running_waits_for_official_api_readiness(self) -> None: |
| server = LMStudioServer(cli_path=__file__) |
| status = {"running": True, "returncode": 0, "stdout": "", "stderr": ""} |
| readiness = iter((False, True)) |
| with patch.object(server, "status", return_value=status), patch.object( |
| server, "_api_ready", side_effect=lambda: next(readiness) |
| ), patch("agent_harness.lm_studio_management.time.sleep"): |
| result = server.ensure_running() |
| self.assertEqual(result["action"], "already_running") |
| self.assertEqual(result["status"], status) |
|
|
| def test_exclusive_switch_unloads_previous_model(self) -> None: |
| manager = FakeResidencyManager() |
| first = manager.ensure_exclusive("embedding", 8192) |
| second = manager.ensure_exclusive("qwen", 65536) |
| self.assertFalse(first.reused) |
| self.assertEqual(second.unloaded_instances, ("embedding",)) |
| self.assertEqual(second.after_instances, ("qwen",)) |
| self.assertEqual(manager.loaded_instances()[0]["model_key"], "qwen") |
|
|
| def test_matching_model_and_context_are_reused(self) -> None: |
| manager = FakeResidencyManager() |
| manager.ensure_exclusive("qwen", 65536) |
| transition = manager.ensure_exclusive("qwen", 65536) |
| self.assertTrue(transition.reused) |
| self.assertEqual(transition.unloaded_instances, ()) |
|
|
| def test_context_change_forces_reload(self) -> None: |
| manager = FakeResidencyManager() |
| manager.ensure_exclusive("qwen", 65536) |
| transition = manager.ensure_exclusive("qwen", 32768) |
| self.assertFalse(transition.reused) |
| self.assertEqual(transition.unloaded_instances, ("qwen",)) |
| self.assertEqual(manager.loaded_instances()[0]["config"]["context_length"], 32768) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
