tests/test_lm_studio.py

from __future__ import annotations

from pathlib import Path
import unittest

from agent_harness.lm_studio import (
    LMStudioClient,
    LMStudioError,
    normalize_identity,
    select_expected_model,
)
from agent_harness.specs import load_models


ROOT = Path(__file__).resolve().parents[1]


class LMStudioIdentityTests(unittest.TestCase):
    def setUp(self) -> None:
        self.spec = load_models(ROOT)["M001"]

    def test_identity_normalization(self) -> None:
        self.assertEqual(normalize_identity("Qwen3.6-35B-A3B"), "qwen3635ba3b")

    def test_resolves_expected_model_and_native_metadata(self) -> None:
        resolved = select_expected_model(
            self.spec,
            [{"id": "qwen/qwen3.6-35b-a3b", "object": "model"}],
            [
                {
                    "key": "qwen/qwen3.6-35b-a3b",
                    "selected_variant": "qwen/qwen3.6-35b-a3b@4bit",
                    "format": "mlx",
                    "quantization": {"name": "4bit"},
                    "loaded_instances": [{"config": {"context_length": 262144}}],
                    "capabilities": {"reasoning": {"default": "on"}},
                }
            ],
        )
        self.assertEqual(resolved.inference_key, "qwen/qwen3.6-35b-a3b")
        self.assertEqual(
            resolved.native_record["selected_variant"],
            "qwen/qwen3.6-35b-a3b@4bit",
        )

    def test_rejects_missing_model(self) -> None:
        with self.assertRaises(LMStudioError):
            select_expected_model(self.spec, [{"id": "some-other-model"}])

    def test_rejects_ambiguous_variants(self) -> None:
        with self.assertRaises(LMStudioError):
            select_expected_model(
                self.spec,
                [
                    {"id": "qwen3.6-35b-a3b-q4"},
                    {"id": "qwen3.6-35b-a3b-q8"},
                ],
            )

    def test_inference_probe_rejects_http_success_without_marker(self) -> None:
        client = LMStudioClient(self.spec)
        client.chat_completions = lambda *args, **kwargs: {  # type: ignore[method-assign]
            "choices": [
                {
                    "finish_reason": "length",
                    "message": {"content": "", "reasoning_content": "unfinished"},
                }
            ]
        }
        with self.assertRaises(LMStudioError):
            client.inference_probe(self.spec.expected_inference_key)

    def test_inference_probe_accepts_exact_visible_marker(self) -> None:
        client = LMStudioClient(self.spec)
        response = {
            "choices": [
                {
                    "finish_reason": "stop",
                    "message": {"content": "\nMODEL_OK", "reasoning_content": "checked"},
                }
            ]
        }
        client.chat_completions = lambda *args, **kwargs: response  # type: ignore[method-assign]
        self.assertIs(client.inference_probe(self.spec.expected_inference_key), response)


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