tests/test_live.py

"""Live tests: cheap real-HF reads. Needs VOICE_TEST_LIVE=1. No big downloads here."""
import json
import struct
import sys
import tempfile
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from _helpers import ABL_FILE, ABL_REPO, BASE_FILE, BASE_REPO, CATGIRL_REPO  # noqa: E402
from _helpers import needs_live, voice  # noqa: E402

import numpy as np


def shard_url(repo, filename):
    return f"https://huggingface.co/{repo}/resolve/main/{filename}"


@needs_live
class TestHeadResolve(unittest.TestCase):
    def test_catgirl_head_resolves(self):
        idx = voice._hf_model_index(CATGIRL_REPO)
        self.assertIsNotNone(idx)
        names = list((idx.get("weight_map") or {}).keys())
        self.assertGreater(len(names), 100)
        found = voice.detect_output_tensor(names, None)
        self.assertIsNotNone(found)
        self.assertTrue(found.endswith(("embed_tokens.weight", "lm_head.weight",
                                        "output.weight")), found)

    def test_catgirl_tied_config(self):
        cfg = voice._hf_config(CATGIRL_REPO)
        self.assertIsNotNone(cfg)
        self.assertTrue(cfg.get("tie_word_embeddings"))

    def test_abliterated_gguf_head_resolves(self):
        src = voice.open_source(shard_url(ABL_REPO, ABL_FILE))
        self.assertEqual(src.kind, "gguf")
        names = src.names()
        # Qwen3 GGUF ties the head: token_embd.weight IS the output projection
        self.assertIn("token_embd.weight", names)
        ref = src.ref("token_embd.weight")
        self.assertEqual(ref.dtype, "Q8_0")
        self.assertEqual(len(ref.shape), 2)


@needs_live
class TestSlices(unittest.TestCase):
    def test_safetensors_head_slice_decodes(self):
        idx = voice._hf_model_index(CATGIRL_REPO)
        wm = idx["weight_map"]
        name = voice.detect_output_tensor(list(wm.keys()), None)
        url = shard_url(CATGIRL_REPO, wm[name])
        raw, st, _ = voice.http_range(url, 0, 7)
        self.assertEqual(st, 206)
        hs = struct.unpack("<Q", raw)[0]
        raw, st, _ = voice.http_range(url, 8, 8 + hs - 1)
        info = json.loads(raw)[name]
        begin, end = info["data_offsets"]
        n = min(1 << 20, end - begin)
        d = Path(tempfile.mkdtemp(prefix="live-slice-"))
        combined = voice.download_range(url, 8 + hs + begin, 8 + hs + begin + n - 1,
                                        d, None, label="live-head-slice", n_parts=2)
        arr = voice.decode_to_f32(combined.read_bytes(), info["dtype"],
                                  (n // voice.dtype_bytes(info["dtype"]),))
        self.assertTrue(bool(np.isfinite(arr.astype(np.float32)).all()))
        self.assertGreater(float(np.abs(arr).max()), 0.0)

    def test_remote_gguf_small_tensor_reads(self):
        # smallest tensor in the file: proves remote GGUF header + ranged read + dequant
        src = voice.open_source(shard_url(ABL_REPO, ABL_FILE))
        refs = [(n, src.ref(n)) for n in src.names()]
        refs = [(n, r) for n, r in refs if len(r.shape) <= 2]
        n, r = min(refs, key=lambda nr: int(np.prod(nr[1].shape)))
        arr = src.read_f32(n)
        self.assertEqual(tuple(arr.shape), tuple(r.shape))
        self.assertTrue(bool(np.isfinite(arr).all()))


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