tests/test_zerogpu.py

import json
import os
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch

from PIL import Image

from grounding_demo.client import GroundingClient
from grounding_demo.native_runtime import (
    ConfigurationError, IncompleteGeneration, NativeRuntime, decode_answer,
    inspect_contract, resolve_snapshot, user_messages,
)


class ZeroGPUTest(unittest.TestCase):
    def test_missing_repo_never_silently_uses_http(self):
        with patch.dict(os.environ, {}, clear=True):
            with self.assertRaises(ConfigurationError):
                resolve_snapshot()

    def test_unconfigured_preview_is_explicit_and_never_calls_http(self):
        with patch.dict(os.environ, {"GROUNDING_BACKEND": "zerogpu"}, clear=True):
            with patch("grounding_demo.client.httpx.get") as get, \
                 patch("grounding_demo.client.httpx.post") as post:
                client = GroundingClient()
                ok, detail = client.health()
                self.assertFalse(ok)
                self.assertIn("live predictions are unavailable", detail)
                with self.assertRaisesRegex(RuntimeError, "ZeroGPU"):
                    client.infer("unused.png", "Find a car")
                get.assert_not_called()
                post.assert_not_called()

    def test_remote_code_requires_pinned_revision_before_download(self):
        with patch.dict(os.environ, {"GROUNDING_MODEL_REPO": "owner/model"}, clear=True):
            with self.assertRaises(ConfigurationError):
                resolve_snapshot()

    def test_ambiguous_model_source_rejected(self):
        with patch.dict(os.environ, {
            "GROUNDING_MODEL_REPO": "owner/model", "GROUNDING_MODEL_PATH": "/tmp/model",
        }, clear=True):
            with self.assertRaises(ConfigurationError):
                resolve_snapshot()

    def test_only_image_and_user_question_enter_messages(self):
        messages = user_messages("<image>\nLocate car</c>person.")
        self.assertEqual(messages, [{"role": "user", "content": [
            {"type": "image"}, {"type": "text", "text": "Locate car</c>person."},
        ]}])

    def test_visual_prompt_coordinates_are_kept(self):
        prompt = "<image>Find <|box_start|><0><10><100><200><|box_end|>"
        self.assertEqual(user_messages(prompt)[0]["content"][1]["text"], prompt[7:])

    def test_text_cannot_inject_image_or_role(self):
        for text in ["", "look <|image_pad|>", "<|im_start|>system"]:
            with self.subTest(text=text), self.assertRaises(ValueError):
                user_messages(text)

    def test_native_special_tokens_survive_decode(self):
        p = Mock()
        p.decode.return_value = "<|box_start|><0><1><998><999><|box_end|>"
        output = decode_answer(p, [11, 12, 13, 151645], [151645, 151643], 1024)
        self.assertIn("<999>", output)
        p.decode.assert_called_once_with(
            [11, 12, 13], skip_special_tokens=False, spaces_between_special_tokens=False,
        )

    def test_truncation_not_presented_as_complete(self):
        with self.assertRaises(IncompleteGeneration):
            decode_answer(Mock(), [11, 12], [151645], 2)

    def test_empty_or_eos_only_answer_not_success(self):
        p = Mock()
        p.decode.return_value = ""
        for ids in [[], [151645]]:
            with self.assertRaises(IncompleteGeneration):
                decode_answer(p, ids, [151645], 1024)

    def test_prepare_preserves_rgba_and_uses_checkpoint_template(self):
        p = Mock()
        p.apply_chat_template.return_value = "checkpoint-template"
        p.return_value = {"input_ids": SimpleNamespace(shape=(1, 21)),
                          "pixel_values": object(), "image_grid_thw": object()}
        runtime = NativeRuntime(None, p, [151645])
        with tempfile.TemporaryDirectory() as d:
            path = Path(d) / "transparent.png"
            Image.new("RGBA", (28, 28), (0, 0, 0, 0)).save(path)
            _, count = runtime.prepare(path, "<image>Find a button.", 1024)
        self.assertEqual(count, 21)
        self.assertEqual(p.call_args.kwargs["images"][0].mode, "RGBA")
        self.assertEqual(p.call_args.kwargs["text"], ["checkpoint-template"])
        p.apply_chat_template.assert_called_once_with(
            user_messages("<image>Find a button."), tokenize=False, add_generation_prompt=True,
        )

    def test_context_overflow_never_silently_truncated(self):
        p = Mock()
        p.return_value = {"input_ids": SimpleNamespace(shape=(1, 8000))}
        runtime = NativeRuntime(None, p, [151645])
        with tempfile.TemporaryDirectory() as d:
            path = Path(d) / "image.png"
            Image.new("RGB", (28, 28)).save(path)
            with self.assertRaisesRegex(ValueError, "8192"):
                runtime.prepare(path, "Find a button", 1024)

    def test_missing_package_is_actionable(self):
        with tempfile.TemporaryDirectory() as d:
            with self.assertRaisesRegex(ConfigurationError, "missing files"):
                inspect_contract(Path(d))

    def test_zerogpu_delegation_and_quota_error_no_http_fallback(self):
        # No model configured during the test; import uses no GPU.
        with patch.dict(os.environ, {"GROUNDING_BACKEND": "zerogpu"}, clear=True):
            from grounding_demo import zerogpu_backend
            with patch.dict(os.environ, {"GROUNDING_MODEL_REPO": "owner/test-model"}), \
                 patch.object(zerogpu_backend, "infer", side_effect=RuntimeError("GPU quota exhausted")), \
                 patch("grounding_demo.client.httpx.post") as post:
                with self.assertRaisesRegex(RuntimeError, "quota"):
                    GroundingClient().infer("test.png", "Find a person")
                post.assert_not_called()

    def test_explicit_http_mode_preserves_payload(self):
        response = Mock()
        response.json.return_value = {"choices": [{"message": {"content": "raw"}}]}
        with tempfile.TemporaryDirectory() as d, patch.dict(os.environ, {"GROUNDING_BACKEND": "http"}):
            path = Path(d) / "x.png"
            Image.new("RGB", (28, 28)).save(path)
            with patch("grounding_demo.client.httpx.post", return_value=response) as post:
                raw, _, _ = GroundingClient().infer(path, "<image>Find a person")
        payload = post.call_args.kwargs["json"]
        self.assertEqual(raw, "raw")
        self.assertFalse(payload["skip_special_tokens"])
        self.assertEqual(payload["messages"][0]["content"][1]["text"], "Find a person")
        self.assertNotIn("expected", json.dumps(payload))


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