lmstudiofixed/tests/test_lmstudio_wrapper_generator.py

import importlib.util
import json
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 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"])


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