claude-mode/tests/python/test_cm_json.py

"""Unit tests for linux/cm-json.py - the JSON engine behind the POSIX CLI.

Standard library only. cm-json.py is loaded from the repository, so it reads the
repository's providers.json (it looks one level above itself). Every test that
writes works in a temporary directory.
"""

import contextlib
import importlib.util
import io
import json
import os
import shutil
import sys
import tempfile
import time
import unittest

REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
_spec = importlib.util.spec_from_file_location("cm_json", os.path.join(REPO, "linux", "cm-json.py"))
cm = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(cm)


def call(fn, argv, stdin=""):
    """Run a cm-json command function; return what it printed."""
    out, old = io.StringIO(), sys.stdin
    sys.stdin = io.StringIO(stdin)
    try:
        with contextlib.redirect_stdout(out):
            fn(argv)
    finally:
        sys.stdin = old
    return out.getvalue()


def call_json(fn, argv, stdin=""):
    return json.loads(call(fn, argv, stdin))


class TempDir(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="cm-json-test.")

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def path(self, *parts):
        return os.path.join(self.tmp, *parts)

    def write_json(self, rel, data):
        p = self.path(rel)
        os.makedirs(os.path.dirname(p), exist_ok=True)
        with open(p, "w", encoding="utf-8") as fh:
            json.dump(data, fh)
        return p


# ---------------------------------------------------------------------------
# providers.json
# ---------------------------------------------------------------------------

class Providers(unittest.TestCase):
    def test_order(self):
        self.assertEqual([p["id"] for p in cm.providers()],
                         ["openrouter", "zai", "lmstudio", "ollama", "custom"])

    def test_builtin_defaults(self):
        self.assertEqual(cm.builtin_default_preset(), {
            "openrouter": "default", "zai": "zai", "lmstudio": "lmstudio",
            "ollama": "ollama", "custom": "custom"})

    def test_tsv_column_contract(self):
        # The bash CLI addresses these columns by number (see prov_field in
        # linux/claude-mode). Reordering them silently breaks it; append only.
        self.assertEqual([name for name, _ in cm.PROVIDER_TSV], [
            "id", "aliases", "title", "label", "color", "defaultPreset",
            "serverEditable", "probe", "probePaths", "catalogueKind", "perServer",
            "setupKey", "setupModels", "keyUrl", "guardrail", "doctor",
            "defaultKeyRef", "literalToken", "defaultBaseUrl", "serverHint", "serverStart"])
        for line in call(cm.cmd_provider_tsv, []).splitlines():
            self.assertEqual(len(line.split("\t")), len(cm.PROVIDER_TSV), line)

    def test_resolve(self):
        for word, want in [("zai", "zai"), ("Z.AI", "zai"), ("z-ai", "zai"),
                           ("lm-studio", "lmstudio"), ("OLLAMA", "ollama")]:
            self.assertEqual(call(cm.cmd_provider_resolve, [word]).strip(), want)
        for word in ("anthropic", "nope", ""):
            with self.assertRaises(SystemExit):
                call(cm.cmd_provider_resolve, [word])

    def test_static_catalogue(self):
        self.assertEqual(call(cm.cmd_provider_static, ["zai"]).splitlines()[0].split("\t")[0], "glm-5.3")


class Scaffold(unittest.TestCase):
    def blank(self, **over):
        base = {"provider": None, "description": "new preset", "baseUrl": None, "auth": None,
                "models": {t: "" for t in cm.TIERS}, "subagentModel": "inherit",
                "gatewayModelDiscovery": False, "contextTokens": None}
        base.update(over)
        return json.dumps(base, indent=2) + "\n"

    # The original three, exactly as the hard-coded scaffold produced them
    # before providers.json existed.
    def test_openrouter(self):
        self.assertEqual(call(cm.cmd_scaffold, ["openrouter"]), self.blank(
            provider="openrouter", baseUrl="https://openrouter.ai/api",
            auth={"mode": "vault", "keyRef": "openrouter"}, gatewayModelDiscovery=True,
            contextTokens=1000000))

    def test_zai(self):
        want = json.loads(self.blank(
            provider="zai", baseUrl="https://api.z.ai/api/anthropic",
            auth={"mode": "vault", "keyRef": "zai"}, contextTokens=1000000))
        want["extraEnv"] = {"API_TIMEOUT_MS": "3000000", "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}
        self.assertEqual(call(cm.cmd_scaffold, ["zai"]), json.dumps(want, indent=2) + "\n")

    def test_lmstudio(self):
        want = json.loads(self.blank(
            provider="lmstudio", baseUrl="http://127.0.0.1:1234",
            auth={"mode": "literal", "token": "lmstudio"}, contextTokens=262144))
        want["extraEnv"] = {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}
        self.assertEqual(call(cm.cmd_scaffold, ["lmstudio"]), json.dumps(want, indent=2) + "\n")

    def test_new_providers(self):
        o = json.loads(call(cm.cmd_scaffold, ["ollama"]))
        self.assertEqual((o["baseUrl"], o["auth"], o["contextTokens"]),
                         ("http://127.0.0.1:11434", {"mode": "literal", "token": "ollama"}, 65536))
        c = json.loads(call(cm.cmd_scaffold, ["custom"]))
        self.assertEqual((c["baseUrl"], c["auth"]), ("", {"mode": "vault", "keyRef": "custom"}))

    def test_unknown(self):
        with self.assertRaises(SystemExit):
            call(cm.cmd_scaffold, ["nope"])


# ---------------------------------------------------------------------------
# Catalogues
# ---------------------------------------------------------------------------

class Parsers(unittest.TestCase):
    def test_openrouter(self):
        src = {"data": [{"id": "b/x", "context_length": 1000, "pricing": {"prompt": "0.000001", "completion": "0.000002"}},
                        {"id": "a/y", "context_length": 5, "pricing": {}}]}
        self.assertEqual(call(cm.cmd_or_models, [], json.dumps(src)).splitlines(),
                         ["a/y\t5\t\t", "b/x\t1000\t1.0\t2.0"])

    def test_lmstudio(self):
        src = {"data": [{"id": "m", "state": "loaded", "max_context_length": 8192}]}
        self.assertEqual(call(cm.cmd_lms_models, [], json.dumps(src)), "m\tloaded\t8192\n")

    def test_ollama_placeholders(self):
        # "-" rather than empty, so bash's tab-IFS read cannot merge columns.
        src = {"models": [{"name": "glm:cloud", "details": {}},
                          {"name": "q:latest", "details": {"parameter_size": "7B", "quantization_level": "Q4"}}]}
        self.assertEqual(call(cm.cmd_ollama_models, [], json.dumps(src)).splitlines(),
                         ["glm:cloud\t-\t-\t-", "q:latest\t7B\tQ4\t-"])

    def test_openai_both_shapes(self):
        self.assertEqual(call(cm.cmd_openai_models, [], json.dumps({"data": [{"id": "b"}, {"id": "a"}]})), "a\nb\n")
        self.assertEqual(call(cm.cmd_openai_models, [], json.dumps([{"id": "x"}])), "x\n")


class Cache(TempDir):
    def cache(self, provider, ok, base, tsv):
        call(cm.cmd_cache_models, [self.path("c.json"), provider, "1" if ok else "0", base], tsv)
        with open(self.path("c.json")) as fh:
            return json.load(fh)["providers"][provider]

    def test_per_server_and_failure(self):
        node = self.cache("ollama", True, "http://a:1/", "q:latest\t7B\tQ4\t-\n")
        self.assertEqual((node["baseUrl"], node["models"]), ("http://a:1", [{"id": "q:latest", "note": "7B Q4"}]))
        node = self.cache("ollama", False, "http://a:1", "")
        self.assertEqual((node["ok"], len(node["models"])), (False, 1))   # a stale list beats none
        node = self.cache("ollama", False, "http://b:2", "")
        self.assertEqual((node["baseUrl"], node["models"]), ("http://b:2", []))   # another server's list is wrong

    def test_openrouter_prices_and_no_secrets(self):
        node = self.cache("openrouter", True, "", "a/y\t5\t0.1\t0.2\n")
        self.assertEqual(node["models"], [{"id": "a/y", "contextTokens": 5, "priceIn": 0.1, "priceOut": 0.2}])
        self.assertNotIn("baseUrl", node)
        with open(self.path("c.json")) as fh:
            text = fh.read()
        for word in ("keyRef", "token", "auth"):
            self.assertNotIn('"%s"' % word, text)


# ---------------------------------------------------------------------------
# Session transcripts: scan, dismiss, repair
# ---------------------------------------------------------------------------

GOOD = {"type": "assistant", "message": {"id": "msg_1", "model": "claude-opus-5", "content": [{"type": "text", "text": "hi"}]}}
FOREIGN = {"type": "assistant", "message": {"id": "gen-1-a", "model": "deepseek/x", "content": [{"type": "text", "text": "yo"}]}}


class Sessions(TempDir):
    def setUp(self):
        super().setUp()
        self.projects = self.path("projects")
        os.makedirs(os.path.join(self.projects, "-p"))
        self.ignored = self.path("ignored.json")

    def transcript(self, sid, lines, age_days=0):
        p = os.path.join(self.projects, "-p", sid + ".jsonl")
        with open(p, "w") as fh:
            fh.write("\n".join(json.dumps(x) for x in lines) + "\n")
        t = time.time() - age_days * 86400
        os.utime(p, (t, t))
        return p

    def scan(self, max_age=""):
        return call_json(cm.cmd_scan_sessions, [self.projects, max_age, self.ignored])

    def test_scan_routes_by_reason(self):
        self.transcript("new", [GOOD, FOREIGN])
        self.transcript("old", [GOOD, FOREIGN], age_days=10)
        self.transcript("healthy", [GOOD])
        self.transcript("native", [FOREIGN])            # ran entirely on a gateway: not damage
        d = self.scan()
        self.assertEqual([b["sessionId"] for b in d["broken"]], ["new"])
        self.assertEqual([(i["sessionId"], i["reason"]) for i in d["ignored"]], [("old", "stale")])
        self.assertEqual(self.scan("0")["count"], 2)

    def test_dismiss_list_prune(self):
        self.transcript("new", [GOOD, FOREIGN])
        gone = self.transcript("gone", [GOOD, FOREIGN])
        call(cm.cmd_ignore_session, [self.ignored, self.projects, "add", "new"])
        call(cm.cmd_ignore_session, [self.ignored, self.projects, "add", "gone"])
        self.assertEqual([(i["sessionId"], i["reason"]) for i in self.scan()["ignored"]],
                         [("gone", "dismissed"), ("new", "dismissed")])
        rows = call_json(cm.cmd_ignore_session, [self.ignored, self.projects, "list"])["sessions"]
        self.assertTrue(all(r["broken"] and r["exists"] for r in rows))
        os.remove(gone)
        out = call_json(cm.cmd_ignore_session, [self.ignored, self.projects, "remove", "new"])
        self.assertEqual((out["changed"], out["pruned"], out["count"]), (True, 1, 0))
        with self.assertRaises(SystemExit):
            call(cm.cmd_ignore_session, [self.ignored, self.projects, "add", "../x"])
        with self.assertRaises(SystemExit):
            call(cm.cmd_ignore_session, [self.ignored, self.projects, "add", "nosuch"])

    def test_repair_truncates_and_hands_back(self):
        p = self.transcript("fix", [GOOD, FOREIGN])
        v = call_json(cm.cmd_repair_session, [p, "--apply"])
        self.assertTrue(v["applied"] and v["reinjected"])
        self.assertTrue(os.path.exists(v["backup"]) and os.path.exists(v["recovered"]))
        with open(p) as fh:
            lines = [json.loads(x) for x in fh.read().splitlines()]
        self.assertEqual(lines[0]["message"]["id"], "msg_1")
        self.assertTrue(lines[1]["isMeta"])
        self.assertIn("<recovered-transcript>", lines[1]["message"]["content"])
        self.assertTrue(call_json(cm.cmd_repair_session, [p])["healthy"])


# ---------------------------------------------------------------------------
# Presets: rename, defaults, auth
# ---------------------------------------------------------------------------

class Presets(TempDir):
    def setUp(self):
        super().setUp()
        self.presets = self.path("presets")
        shutil.copytree(os.path.join(REPO, "presets"), self.presets)
        self.state = self.write_json("state.json", {"mode": "openrouter", "preset": "default", "writtenEnvKeys": []})
        self.defaults = self.write_json("defaults.json", {"openrouter": "default"})

    def load(self, p):
        with open(p) as fh:
            return json.load(fh)

    def test_rename_repoints_state_and_default(self):
        out = call_json(cm.cmd_preset_rename, [self.presets, "default", "daily", self.state, self.defaults])
        self.assertEqual(out, {"renamed": True, "active": True, "default": True})
        self.assertEqual(self.load(self.state)["preset"], "daily")
        self.assertEqual(self.load(self.defaults)["openrouter"], "daily")
        self.assertFalse(os.path.exists(os.path.join(self.presets, "default.json")))
        with self.assertRaises(SystemExit):
            call(cm.cmd_preset_rename, [self.presets, "daily", "zai", self.state])
        with self.assertRaises(SystemExit):
            call(cm.cmd_preset_rename, [self.presets, "nosuch", "x", self.state])

    def test_set_default_and_clear(self):
        call(cm.cmd_set_default, [self.defaults, "zai", "zai"])
        call(cm.cmd_set_default, [self.defaults, "openrouter", ""])
        self.assertEqual(self.load(self.defaults), {"zai": "zai"})

    def test_auth_of(self):
        self.assertEqual(call(cm.cmd_auth_of, [os.path.join(self.presets, "zai.json")]), "vault\tzai\n")
        self.assertEqual(call(cm.cmd_auth_of, [os.path.join(self.presets, "ollama.json")]), "literal\topenrouter\n")
        with self.assertRaises(SystemExit) as e:
            call(cm.cmd_auth_of, [os.path.join(self.presets, "nosuch.json")])
        self.assertEqual(e.exception.code, 1)

    def test_set_auth_uses_the_providers_names(self):
        p = os.path.join(self.presets, "ollama.json")
        call(cm.cmd_set_auth, [p, "key"])
        self.assertEqual(self.load(p)["auth"], {"mode": "vault", "keyRef": "ollama"})
        call(cm.cmd_set_auth, [p, "none"])
        self.assertEqual(self.load(p)["auth"], {"mode": "literal", "token": "ollama"})


class Health(Presets):
    def health(self, mode="anthropic", preset=""):
        call(cm.cmd_health, [self.tmp, "/nonexistent", mode, preset, "", "file", "9.9.9"])
        return self.load(self.path("health.json"))

    def test_effective_defaults(self):
        self.assertEqual(self.health()["defaultPresetFor"]["openrouter"], "default")
        shutil.copy(os.path.join(self.presets, "default.json"), os.path.join(self.presets, "cheap.json"))
        self.write_json("defaults.json", {"openrouter": "cheap"})
        h = self.health()
        self.assertEqual((h["defaultPresetFor"]["openrouter"], h["defaultPresetChosen"]), ("cheap", {"openrouter": "cheap"}))
        os.remove(os.path.join(self.presets, "cheap.json"))           # a vanished choice falls back
        self.assertEqual(self.health()["defaultPresetFor"]["openrouter"], "default")
        os.remove(os.path.join(self.presets, "default.json"))         # and with the built-in gone, the first by name
        shutil.copy(os.path.join(self.presets, "zai.json"), os.path.join(self.presets, "alpha.json"))
        alpha = self.load(os.path.join(self.presets, "alpha.json"))
        alpha["provider"] = "openrouter"
        self.write_json("presets/alpha.json", alpha)
        self.assertEqual(self.health()["defaultPresetFor"]["openrouter"], "alpha")

    def test_publishes_providers(self):
        h = self.health()
        self.assertEqual(h["version"], "9.9.9")
        byid = {p["id"]: p for p in h["providers"]}
        self.assertEqual(list(byid), ["openrouter", "zai", "lmstudio", "ollama", "custom"])
        self.assertTrue(byid["ollama"]["serverEditable"] and byid["ollama"]["perServerCatalogue"])
        self.assertFalse(byid["zai"]["serverEditable"])
        self.assertEqual(byid["custom"]["defaultBaseUrl"], "")
        self.assertTrue(byid["openrouter"]["logo"])


class OllamaCtx(unittest.TestCase):
    def test_show(self):
        src = {"model_info": {"general.architecture": "q", "q.context_length": 262144}}
        self.assertEqual(call(cm.cmd_ollama_ctx, ["show"], json.dumps(src)), "262144\n")

    def test_ps_matches_latest(self):
        src = {"models": [{"name": "qwen3-coder:latest", "model": "qwen3-coder:latest", "context_length": 4096}]}
        self.assertEqual(call(cm.cmd_ollama_ctx, ["ps", "qwen3-coder"], json.dumps(src)), "4096\n")
        self.assertEqual(call(cm.cmd_ollama_ctx, ["ps", "other"], json.dumps(src)), "")


# ---------------------------------------------------------------------------
# apply - the write into settings.json
# ---------------------------------------------------------------------------

class Apply(TempDir):
    def preset(self, **over):
        base = {"provider": "openrouter", "baseUrl": "https://gw", "auth": {"mode": "vault", "keyRef": "openrouter"},
                "models": {"opus": "x/opus", "sonnet": "x/sonnet"}, "contextTokens": 1000}
        base.update(over)
        return self.write_json("p.json", base)

    def apply(self, mode, preset="", helper="/opt/key helper/h.sh"):
        return call(cm.cmd_apply, [self.path("settings.json"), self.path("state.json"), mode, preset, helper])

    def settings(self):
        with open(self.path("settings.json")) as fh:
            return json.load(fh)

    def test_writes_env_and_quotes_the_helper(self):
        self.apply("openrouter", self.preset())
        s = self.settings()
        self.assertEqual(s["env"]["ANTHROPIC_BASE_URL"], "https://gw")
        self.assertEqual(s["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"], "x/opus")
        self.assertEqual(s["env"]["CLAUDE_CODE_MAX_CONTEXT_TOKENS"], "1000")
        self.assertEqual(s["env"]["ANTHROPIC_API_KEY"], "")
        self.assertEqual(s["apiKeyHelper"], "'/opt/key helper/h.sh'")

    def test_cost_guard(self):
        with self.assertRaises(SystemExit) as e:
            self.apply("openrouter", self.preset(models={"opus": "anthropic/claude-opus-5"}))
        self.assertIn("refusing to switch", str(e.exception))
        self.apply("openrouter", self.preset(models={"opus": "anthropic/claude-opus-5"}, allowAnthropicModels=True))

    def test_extra_env_is_removed_on_the_way_out(self):
        self.apply("openrouter", self.preset(extraEnv={"SOME_CUSTOM_KEY": "1"}))
        self.assertIn("SOME_CUSTOM_KEY", self.settings()["env"])
        self.apply("anthropic")
        s = self.settings()
        self.assertNotIn("env", s)
        self.assertNotIn("apiKeyHelper", s)

    def test_literal_token_and_provider_check(self):
        self.apply("ollama", self.preset(provider="ollama", auth={"mode": "literal", "token": "ollama"}))
        self.assertEqual(self.settings()["env"]["ANTHROPIC_AUTH_TOKEN"], "ollama")
        with self.assertRaises(SystemExit):
            self.apply("zai", self.preset())


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