lmstudiofixed/tests/test_lmstudio_wrapper_generator.py

import importlib.util
import json
import struct
import tempfile
import unittest
from pathlib import Path

MODULE_PATH = (
    Path(__file__).resolve().parent.parent
    / "tools"
    / "lmstudio"
    / "wrapper_generator.py"
)
if MODULE_PATH.exists():
    spec = importlib.util.spec_from_file_location("wrapper_generator", MODULE_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec is not None and spec.loader is not None
    spec.loader.exec_module(module)
    gen = module
else:  # pragma: no cover
    gen = None


class WrapperGeneratorTests(unittest.TestCase):
    def setUp(self):
        if gen is None:
            self.fail("tools.lmstudio.wrapper_generator is not available")

        self.official_model_yaml = """\
model: qwen/qwen3.6-35b-a3b
base:
  - key: lmstudio-community/qwen3.6-35b-a3b-gguf
metadataOverrides:
  domain: llm
  architectures:
    - qwen35moe
  compatibilityTypes:
    - gguf
  reasoning: true
config:
  operation:
    fields:
      - key: llm.prediction.promptTemplate
        value:
          type: jinja
          jinjaPromptTemplate:
            template: |
              {%- if add_generation_prompt %}
                  {%- if enable_thinking is defined and enable_thinking is false %}
                      {{- '<think>\\n\\n</think>\\n\\n' }}
                  {%- else %}
                      {{- '<think>\\n' }}
                  {%- endif %}
              {%- endif %}
customFields:
  - key: enableThinking
    effects:
      - type: setJinjaVariable
        variable: enable_thinking
  - key: preserveThinking
    effects:
      - type: setJinjaVariable
        variable: preserve_thinking
"""
        self.official_manifest = {
            "type": "model",
            "owner": "qwen",
            "name": "qwen3.6-35b-a3b",
            "dependencies": [
                {
                    "type": "model",
                    "purpose": "baseModel",
                    "modelKeys": ["lmstudio-community/qwen3.6-35b-a3b-gguf"],
                }
            ],
            "revision": 1,
        }

    def _write_gguf_metadata(self, path: Path, metadata: dict[str, object]) -> None:
        type_string = 8
        type_bool = 7

        def write_string(value: str) -> bytes:
            raw = value.encode("utf-8")
            return struct.pack("<Q", len(raw)) + raw

        def write_value(value: object) -> tuple[int, bytes]:
            if isinstance(value, bool):
                return type_bool, struct.pack("<?", value)
            if isinstance(value, str):
                return type_string, write_string(value)
            raise TypeError(f"Unsupported GGUF test metadata value: {value!r}")

        payload = bytearray()
        payload.extend(b"GGUF")
        payload.extend(struct.pack("<I", 3))
        payload.extend(struct.pack("<Q", 0))
        payload.extend(struct.pack("<Q", len(metadata)))

        for key, value in metadata.items():
            value_type, raw_value = write_value(value)
            payload.extend(write_string(key))
            payload.extend(struct.pack("<I", value_type))
            payload.extend(raw_value)

        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(payload)

    def test_build_wrapper_files_contains_reasoning_and_target_key(self):
        result = gen.build_wrapper_files(
            target_model_key="qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
            wrapper_owner="local",
            wrapper_name="qwen3.6-35b-a3b-hauhaucs-aggressive",
            official_model_yaml_text=self.official_model_yaml,
            official_manifest=self.official_manifest,
        )

        self.assertIn("reasoning: true", result["model.yaml"])
        self.assertIn("key: enableThinking", result["model.yaml"])
        self.assertIn("key: preserveThinking", result["model.yaml"])
        self.assertIn(
            "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
            result["model.yaml"],
        )

        manifest = json.loads(result["manifest.json"])
        self.assertEqual(manifest["owner"], "local")
        self.assertEqual(manifest["name"], "qwen3.6-35b-a3b-hauhaucs-aggressive")
        self.assertEqual(
            manifest["dependencies"][0]["modelKeys"],
            ["qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive"],
        )

    def test_dry_run_does_not_write_output(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            output_dir = Path(tmpdir) / "wrapper"
            summary = gen.generate_wrapper(
                output_dir=output_dir,
                target_model_key="qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive",
                wrapper_owner="local",
                wrapper_name="qwen3.6-35b-a3b-hauhaucs-aggressive",
                official_model_yaml_text=self.official_model_yaml,
                official_manifest=self.official_manifest,
                dry_run=True,
            )

            self.assertFalse(output_dir.exists())
            self.assertTrue(summary["dry_run"])
            self.assertIn("model.yaml", summary["files"])
            self.assertIn("manifest.json", summary["files"])
            self.assertIn("README.md", summary["files"])

    def test_collect_model_jobs_skips_lmstudio_community_and_matches_template(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            hub_root = root / "hub" / "models"
            output_root = root / "output"

            self._write_gguf_metadata(
                models_root / "alice" / "qwen-local" / "model.gguf",
                {
                    "general.architecture": "qwen3",
                    "general.name": "Qwen Local",
                    "capabilities.reasoning": True,
                    "capabilities.tools": True,
                },
            )
            self._write_gguf_metadata(
                models_root / "lmstudio-community" / "official" / "model.gguf",
                {
                    "general.architecture": "qwen3",
                    "general.name": "Skip Me",
                },
            )

            template_dir = hub_root / "qwen" / "qwen3.6-35b-a3b"
            template_dir.mkdir(parents=True, exist_ok=True)
            (template_dir / "model.yaml").write_text(
                self.official_model_yaml, encoding="utf-8"
            )
            (template_dir / "manifest.json").write_text(
                json.dumps(self.official_manifest, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )

            summary = gen.generate_wrappers_for_models(
                models_root=models_root,
                hub_root=hub_root,
                output_root=output_root,
                dry_run=True,
            )

            self.assertTrue(summary["dry_run"])
            self.assertEqual(len(summary["generated"]), 1)
            self.assertEqual(len(summary["skipped"]), 1)

            generated = summary["generated"][0]
            self.assertEqual(generated["owner"], "alice")
            self.assertEqual(generated["name"], "qwen-local")
            self.assertEqual(generated["model_key"], "model")
            self.assertIn("qwen3.6-35b-a3b", generated["template_path"])
            self.assertTrue(generated["capabilities"]["reasoning"])
            self.assertTrue(generated["capabilities"]["tools"])
            self.assertFalse((output_root / "alice" / "qwen-local").exists())

            skipped = summary["skipped"][0]
            self.assertEqual(skipped["reason"], "lmstudio-community")

    def test_generate_wrappers_writes_files_and_merges_capabilities(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            hub_root = root / "hub" / "models"
            output_root = root / "hub" / "models"

            self._write_gguf_metadata(
                models_root / "alice" / "qwen-local" / "model.gguf",
                {
                    "general.architecture": "qwen3",
                    "general.name": "Qwen Local",
                    "capabilities.reasoning": False,
                },
            )

            template_dir = hub_root / "qwen" / "qwen3.6-35b-a3b"
            template_dir.mkdir(parents=True, exist_ok=True)
            (template_dir / "model.yaml").write_text(
                self.official_model_yaml, encoding="utf-8"
            )
            (template_dir / "manifest.json").write_text(
                json.dumps(self.official_manifest, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )

            summary = gen.generate_wrappers_for_models(
                models_root=models_root,
                hub_root=hub_root,
                output_root=output_root,
                dry_run=False,
            )

            self.assertEqual(len(summary["generated"]), 1)
            output_dir = output_root / "alice" / "qwen-local"
            self.assertTrue(output_dir.exists())

            model_yaml = (output_dir / "model.yaml").read_text(encoding="utf-8")
            manifest = json.loads(
                (output_dir / "manifest.json").read_text(encoding="utf-8")
            )

            self.assertIn("model: alice/qwen-local", model_yaml)
            self.assertIn('base: "model"', model_yaml)
            self.assertIn("reasoning: false", model_yaml)
            self.assertIn("key: enableThinking", model_yaml)
            self.assertEqual(manifest["owner"], "alice")
            self.assertEqual(manifest["name"], "qwen-local")
            self.assertEqual(manifest["dependencies"][0]["modelKeys"], ["model"])


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