clawdmeter/daemon/tests/test_zai.py

"""v3 M3 — z.ai (GLM) provider: usage from the monitor/quota endpoint.

Uses httpx.MockTransport so no real z.ai call is made. Covers the token-window
mapping (5h -> s, weekly -> w), raw-key auth, the ignored monthly TIME_LIMIT,
and auth/error handling.
"""

import asyncio
import json
import time

import httpx

from daemon.providers.zai import ZaiProvider, DEFAULT_HOST, _MONITOR_PATH


def _run(coro):
    # test_zai sorts last; an earlier module's asyncio.run() can leave the shared
    # loop closed, so fall back to a fresh loop rather than raising (Python 3.13).
    try:
        loop = asyncio.get_event_loop()
        if loop.is_closed():
            raise RuntimeError
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    return loop.run_until_complete(coro)


def _config(tmp_path, monkeypatch, zai: dict):
    p = tmp_path / "config.json"
    p.write_text(json.dumps({"version": 2, "providers": {"zai": zai}}), encoding="utf-8")
    monkeypatch.setenv("CLAWDMETER_CONFIG", str(p))
    monkeypatch.setenv("LOCALAPPDATA", str(tmp_path))


def _tok(pct, reset_ms, unit=3, number=5):
    return {"type": "TOKENS_LIMIT", "unit": unit, "number": number,
            "percentage": pct, "nextResetTime": reset_ms}


def _time_limit(pct=0):
    return {"type": "TIME_LIMIT", "unit": 5, "number": 1, "usage": 100,
            "currentValue": 0, "remaining": 100, "percentage": pct,
            "nextResetTime": 9999999999999, "usageDetails": []}


def _ok_resp(limits, level="lite"):
    return httpx.Response(200, json={"code": 200, "msg": "Operation successful",
                                     "success": True, "data": {"level": level, "limits": limits}})


def _provider_returning(resp_or_fn):
    def handler(request):
        return resp_or_fn(request) if callable(resp_or_fn) else resp_or_fn
    return ZaiProvider(transport=httpx.MockTransport(handler))


def test_maps_token_windows(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "sk-zai",
                                    "base_url": "https://api.z.ai/api/anthropic"})
    now_ms = time.time() * 1000
    seen = {}

    def handler(request: httpx.Request) -> httpx.Response:
        seen["url"] = str(request.url)
        seen["auth"] = request.headers.get("authorization")
        seen["lang"] = request.headers.get("accept-language")
        seen["method"] = request.method
        return _ok_resp([
            _time_limit(),                                     # ignored monthly web-tool quota
            _tok(1, now_ms + 3600_000, unit=3, number=5),      # 5h window  -> short
            _tok(2, now_ms + 6 * 86400_000, unit=6, number=1),  # weekly window -> weekly
        ])

    st = _run(ZaiProvider(transport=httpx.MockTransport(handler)).poll())
    assert st.ok is True
    assert st.s == 1.0 and 58 <= st.sr <= 60
    assert st.w == 2.0 and 8000 <= st.wr <= 8700    # ~6 days in minutes
    assert st.st == "Lite"
    # raw key auth (no "Bearer"), monitor endpoint on the configured host, GET
    assert seen["method"] == "GET"
    assert seen["url"] == f"{DEFAULT_HOST}{_MONITOR_PATH}"
    assert seen["auth"] == "sk-zai"
    assert seen["lang"] == "en-US,en"


def test_no_key(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": ""})
    calls = {"n": 0}

    def handler(request):
        calls["n"] += 1
        return _ok_resp([])

    st = _run(ZaiProvider(transport=httpx.MockTransport(handler)).poll())
    assert st.ok is False and st.st == "no key" and st.auth_problem is True
    assert calls["n"] == 0


def test_http_401_is_bad_key(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "bad"})
    p = _provider_returning(httpx.Response(401, json={"error": "unauthorized"}))
    st = _run(p.poll())
    assert st.ok is False and st.st == "bad key" and st.auth_problem is True


def test_body_success_false_is_bad_key(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "bad"})
    p = _provider_returning(httpx.Response(200, json={"code": 401, "msg": "auth", "success": False, "data": None}))
    st = _run(p.poll())
    assert st.ok is False and st.st == "bad key" and st.auth_problem is True


def test_single_token_window(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "sk"})
    now_ms = time.time() * 1000
    p = _provider_returning(_ok_resp([_tok(37, now_ms + 3600_000)]))
    st = _run(p.poll())
    assert st.s == 37.0 and st.w == 0.0 and st.wr == -1 and st.ok is True


def test_hundred_percent_is_limited(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "sk"})
    now_ms = time.time() * 1000
    p = _provider_returning(_ok_resp([
        _tok(100, now_ms + 60_000, unit=3, number=5),
        _tok(40, now_ms + 6 * 86400_000, unit=6, number=1)]))
    st = _run(p.poll())
    assert st.st == "limited" and st.s == 100.0


def test_only_time_limit_reports_connected(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "sk"})
    p = _provider_returning(_ok_resp([_time_limit(5)], level="pro"))
    st = _run(p.poll())
    assert st.ok is True and st.st == "Pro" and st.s == 0.0


def test_transient_error_reported(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "sk"})

    def handler(request):
        raise httpx.ConnectError("boom")

    st = _run(ZaiProvider(transport=httpx.MockTransport(handler)).poll())
    assert st.ok is False and st.st == "offline"


def test_custom_base_host(tmp_path, monkeypatch):
    _config(tmp_path, monkeypatch, {"enabled": True, "api_key": "sk",
                                    "base_url": "https://zzz.example/api/anthropic"})
    seen = {}

    def handler(request):
        seen["url"] = str(request.url)
        return _ok_resp([])

    _run(ZaiProvider(transport=httpx.MockTransport(handler)).poll())
    assert seen["url"] == f"https://zzz.example{_MONITOR_PATH}"   # monitor path on the same host