tests/test_mock_mvp.py

"""Smoke tests for the initial mock MVP."""

from __future__ import annotations

import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import patch

import src.models.llama_cpp_runner as llama_cpp_runner
from src.example_cache import load_sample_generation, sample_trace_path
from src.examples import EXAMPLE_OBJECTS, gradio_examples
from src.models.llama_cpp_runner import (
    generate_diary,
    generate_persona,
    reply_as_object,
    reset_text_runtime_fallbacks,
)
from src.models.vision_runner import understand_object, understand_object_with_metadata
from src.models.vision_runner import probe_vision_runtime
from src.pipeline import generate_object_diary
from src.renderer.share_card import render_share_card
from src.ui.layout import vision_runtime_probe
from src.traces.anonymizer import anonymize_text
from src.traces.logger import build_trace, save_trace
from scripts.generate_sample_traces import generate_sample_traces
from scripts.check_initial_stage import run_checks
from src.config import get_runtime_settings, runtime_status


class FakeMiniCpmModel:
    def __init__(self, response: str) -> None:
        self.response = response

    def chat(self, **_: object) -> str:
        return self.response


class FakeLlamaModel:
    def __init__(self, responses: list[str]) -> None:
        self.responses = responses
        self.calls = 0

    def create_chat_completion(self, **_: object) -> dict:
        self.calls += 1
        response = self.responses.pop(0)
        return {"choices": [{"message": {"content": response}}]}


class MockMvpTest(unittest.TestCase):
    def tearDown(self) -> None:
        reset_text_runtime_fallbacks()

    def test_runtime_defaults_to_mock(self) -> None:
        settings = get_runtime_settings({})
        status = runtime_status(settings)

        self.assertEqual(settings.vision_backend, "mock")
        self.assertEqual(settings.text_backend, "mock")
        self.assertEqual(status["vision"], "mock object understanding")
        self.assertEqual(status["runtime"], "no llama.cpp model connected yet")

    def test_llama_cpp_runtime_status_does_not_expose_model_path(self) -> None:
        status = runtime_status(
            get_runtime_settings(
                {
                    "OBJECTVERSE_TEXT_BACKEND": "llama-cpp",
                    "TEXT_MODEL_PATH": "/Users/leo/private/model.gguf",
                }
            )
        )

        self.assertEqual(status["text"], "llama-cpp text generation")
        self.assertIn("[configured external GGUF]", status["runtime"])
        self.assertNotIn("/Users/leo", status["runtime"])

    def test_llama_cpp_hub_runtime_status_uses_public_repo_summary(self) -> None:
        settings = get_runtime_settings(
            {
                "OBJECTVERSE_TEXT_BACKEND": "llama-cpp",
                "TEXT_MODEL_REPO_ID": "qqyule/objectverse-diary-qwen15b-lora",
                "TEXT_MODEL_FILENAME": "objectverse-diary-qwen15b-lora-v2-q4_k_m.gguf",
            }
        )

        status = runtime_status(settings)

        self.assertEqual(settings.text_model_repo_id, "qqyule/objectverse-diary-qwen15b-lora")
        self.assertEqual(settings.text_model_filename, "objectverse-diary-qwen15b-lora-v2-q4_k_m.gguf")
        self.assertIn("Hub GGUF", status["runtime"])
        self.assertIn("qqyule/objectverse-diary-qwen15b-lora", status["runtime"])
        self.assertNotIn("/home", status["runtime"])
        self.assertNotIn("/Users", status["runtime"])

    def test_llama_cpp_loads_model_from_hub_config_when_path_is_missing(self) -> None:
        previous_model = llama_cpp_runner._LLAMA_MODEL
        previous_path = llama_cpp_runner._LLAMA_MODEL_PATH
        llama_cpp_runner._LLAMA_MODEL = None
        llama_cpp_runner._LLAMA_MODEL_PATH = None

        loaded_paths: list[str] = []

        class FakeLlama:
            def __init__(self, *, model_path: str, **_: object) -> None:
                loaded_paths.append(model_path)

        fake_module = types.ModuleType("llama_cpp")
        fake_module.Llama = FakeLlama

        try:
            with tempfile.TemporaryDirectory() as tmp_dir:
                model_path = Path(tmp_dir) / "model.gguf"
                model_path.write_bytes(b"GGUF")
                settings = get_runtime_settings(
                    {
                        "OBJECTVERSE_TEXT_BACKEND": "llama-cpp",
                        "TEXT_MODEL_REPO_ID": "qqyule/objectverse-diary-qwen15b-lora",
                        "TEXT_MODEL_FILENAME": "objectverse-diary-qwen15b-lora-v2-q4_k_m.gguf",
                    }
                )

                with (
                    patch.dict(sys.modules, {"llama_cpp": fake_module}),
                    patch("src.models.llama_cpp_runner._download_hf_gguf", return_value=str(model_path)),
                ):
                    llama_cpp_runner._load_llama_model("", settings=settings)

            self.assertEqual(loaded_paths, [str(model_path)])
        finally:
            llama_cpp_runner._LLAMA_MODEL = previous_model
            llama_cpp_runner._LLAMA_MODEL_PATH = previous_path

    def test_examples_cover_six_objects(self) -> None:
        self.assertEqual(len(EXAMPLE_OBJECTS), 6)
        self.assertEqual(len(gradio_examples()), 6)
        self.assertTrue(all(len(example) == 1 for example in gradio_examples()))

    def test_sample_generation_cache_loads_committed_example_trace(self) -> None:
        path = sample_trace_path(0)
        result = load_sample_generation(0)

        self.assertIsNotNone(path)
        self.assertIsNotNone(result)
        assert result is not None
        self.assertEqual(result.trace.trace_id, "sample-01")
        self.assertEqual(result.object_understanding.object.name, "coffee mug")
        self.assertEqual(result.trace_path, str(path))

    def test_mock_generation_flow(self) -> None:
        object_understanding = understand_object(
            None,
            "old white coffee mug on my developer desk",
        )
        persona = generate_persona(object_understanding, "Cynical")
        diary = generate_diary(persona, "Cynical")
        share_card = render_share_card(persona, diary)

        self.assertEqual(object_understanding.object.name, "coffee mug")
        self.assertEqual(len(persona.persona.tags), 3)
        self.assertIn("Secret Diary", diary.title)
        self.assertIn("今天", diary.chinese)
        self.assertIn("objectverse-card", share_card)

    def test_llama_cpp_persona_diary_and_chat_accept_valid_json(self) -> None:
        env = {
            "OBJECTVERSE_TEXT_BACKEND": "llama-cpp",
            "TEXT_MODEL_PATH": "/tmp/objectverse-text-model.gguf",
        }
        fake_llama = FakeLlamaModel(
            [
                """
                {"persona":{"object_name":"coffee mug","character_name":"Mugworth","mood":"dry and suspicious","secret_fear":"being left empty forever","core_memory":"It remembers every late-night refill.","complaint":"I am treated like a ceramic fuel tank.","tags":["desk witness","warm archive","quiet judgment"]}}
                """,
                """
                {"title":"Secret Diary - Day 418","english":"Today I held another bitter storm and called it service.","chinese":"今天我又装下一场苦涩风暴，并被称为有用。"}
                """,
                """
                {"reply":"Mugworth: I have seen your deadlines dissolve into coffee rings."}
                """,
            ]
        )

        with (
            patch.dict("os.environ", env, clear=False),
            patch("src.models.llama_cpp_runner._load_llama_model", return_value=fake_llama),
        ):
            object_understanding = understand_object(None, "white coffee mug")
            persona = generate_persona(object_understanding, "Cynical")
            diary = generate_diary(persona, "Cynical")
            reply = reply_as_object(persona.model_dump(mode="json"), "What did you see?")

        self.assertEqual(persona.persona.character_name, "Mugworth")
        self.assertEqual(diary.title, "Secret Diary - Day 418")
        self.assertIn("Mugworth", reply)

    def test_llama_cpp_missing_model_path_falls_back_to_mock(self) -> None:
        env = {"OBJECTVERSE_TEXT_BACKEND": "llama-cpp", "TEXT_MODEL_PATH": ""}

        with patch.dict("os.environ", env, clear=False):
            result = generate_object_diary(None, "dusty black keyboard", "Philosopher", save=False)

        self.assertEqual(result.persona.persona.object_name, "keyboard")
        self.assertIn("text-fallback-to-mock", result.trace.fallbacks)
        self.assertIn("mock-vision-runtime", result.trace.fallbacks)
        self.assertNotIn("mock-text-runtime", result.trace.fallbacks)

    def test_llama_cpp_import_failure_falls_back_to_mock(self) -> None:
        env = {
            "OBJECTVERSE_TEXT_BACKEND": "llama_cpp",
            "TEXT_MODEL_PATH": "/tmp/objectverse-text-model.gguf",
        }

        with (
            patch.dict("os.environ", env, clear=False),
            patch("src.models.llama_cpp_runner._load_llama_model", side_effect=ImportError("no llama_cpp")),
        ):
            result = generate_object_diary(None, "old white coffee mug", "Cynical", save=False)

        self.assertEqual(result.persona.persona.object_name, "coffee mug")
        self.assertIn("text-fallback-to-mock", result.trace.fallbacks)

    def test_llama_cpp_invalid_json_falls_back_to_mock(self) -> None:
        env = {
            "OBJECTVERSE_TEXT_BACKEND": "llama-cpp",
            "TEXT_MODEL_PATH": "/tmp/objectverse-text-model.gguf",
        }

        with (
            patch.dict("os.environ", env, clear=False),
            patch("src.models.llama_cpp_runner._load_llama_model", return_value=FakeLlamaModel(["not json"])),
        ):
            result = generate_object_diary(None, "old white coffee mug", "Cynical", save=False)

        self.assertEqual(result.persona.persona.object_name, "coffee mug")
        self.assertIn("text-fallback-to-mock", result.trace.fallbacks)
        self.assertEqual(result.trace.model_runtime["text"], "llama-cpp text generation")

    def test_pipeline_uses_combined_llama_cpp_persona_and_diary(self) -> None:
        env = {
            "OBJECTVERSE_TEXT_BACKEND": "llama-cpp",
            "TEXT_MODEL_PATH": "/tmp/objectverse-text-model.gguf",
        }
        fake_llama = FakeLlamaModel(
            [
                """
                {
                  "persona": {
                    "object_name": "coffee mug",
                    "character_name": "Mugworth",
                    "mood": "dry and suspicious",
                    "secret_fear": "being left empty forever",
                    "core_memory": "It remembers every late-night refill.",
                    "complaint": "I am treated like a ceramic fuel tank.",
                    "tags": ["desk witness", "warm archive", "quiet judgment"]
                  },
                  "diary": {
                    "title": "Secret Diary - Day 418",
                    "english": "Today I held another bitter storm and called it service.",
                    "chinese": "今天我又装下一场苦涩风暴，并被称为有用。"
                  }
                }
                """,
            ]
        )

        with (
            patch.dict("os.environ", env, clear=False),
            patch("src.models.llama_cpp_runner._load_llama_model", return_value=fake_llama),
        ):
            result = generate_object_diary(None, "old white coffee mug", "Cynical", save=False)

        self.assertEqual(result.persona.persona.character_name, "Mugworth")
        self.assertEqual(result.diary.title, "Secret Diary - Day 418")
        self.assertEqual(fake_llama.calls, 1)
        self.assertNotIn("text-fallback-to-mock", result.trace.fallbacks)

    def test_minicpm_vision_backend_accepts_valid_json(self) -> None:
        response = """
        {"object":{"name":"coffee mug","visible_features":["white ceramic","round handle","desk shadow"],"likely_context":"work desk","confidence":0.88}}
        """
        settings = get_runtime_settings(
            {
                "OBJECTVERSE_VISION_BACKEND": "minicpm-v",
                "VISION_MODEL_ID": "openbmb/MiniCPM-V-2_6",
                "OBJECTVERSE_TEXT_BACKEND": "mock",
            }
        )

        with (
            patch("src.models.vision_runner._load_rgb_image", return_value=object()),
            patch("src.models.vision_runner._load_minicpm_components", return_value=(FakeMiniCpmModel(response), object())),
        ):
            result = understand_object_with_metadata("/tmp/mug.png", "white mug", settings=settings)

        self.assertEqual(result.object_understanding.object.name, "coffee mug")
        self.assertEqual(result.object_understanding.object.confidence, 0.88)
        self.assertEqual(result.fallbacks, [])

    def test_minicpm_vision_backend_falls_back_on_invalid_json(self) -> None:
        settings = get_runtime_settings(
            {
                "OBJECTVERSE_VISION_BACKEND": "minicpm-v",
                "VISION_MODEL_ID": "openbmb/MiniCPM-V-2_6",
                "OBJECTVERSE_TEXT_BACKEND": "mock",
            }
        )

        with (
            patch("src.models.vision_runner._load_rgb_image", return_value=object()),
            patch("src.models.vision_runner._load_minicpm_components", return_value=(FakeMiniCpmModel("not json"), object())),
        ):
            result = understand_object_with_metadata("/tmp/keyboard.png", "dusty black keyboard", settings=settings)

        self.assertEqual(result.object_understanding.object.name, "keyboard")
        self.assertEqual(result.fallbacks, ["vision-fallback-to-mock"])

    def test_vision_runtime_probe_redacts_sensitive_error_markers(self) -> None:
        settings = get_runtime_settings(
            {
                "OBJECTVERSE_VISION_BACKEND": "minicpm-v",
                "VISION_MODEL_ID": "openbmb/MiniCPM-V-2_6",
            }
        )

        with patch(
            "src.models.vision_runner._load_minicpm_components",
            side_effect=RuntimeError("failed with token hf_forbidden in /Users/leo/.env"),
        ):
            probe = probe_vision_runtime(settings=settings, load_model=True)

        serialized = json.dumps(probe, ensure_ascii=False)
        self.assertTrue(probe["minicpm_load_attempted"])
        self.assertFalse(probe["minicpm_load_ok"])
        self.assertNotIn("hf_", serialized)
        self.assertNotIn("HF_TOKEN", serialized)
        self.assertNotIn("/Users/leo", serialized)
        self.assertNotIn(".env", serialized)

    def test_hidden_vision_runtime_probe_returns_safe_json(self) -> None:
        probe = vision_runtime_probe()
        serialized = json.dumps(probe, ensure_ascii=False)

        self.assertIn("backend", probe)
        self.assertIn("torch_import", probe)
        self.assertNotIn("hf_", serialized)
        self.assertNotIn("HF_TOKEN", serialized)

    def test_pipeline_saves_generation_result(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            result = generate_object_diary(
                None,
                "old white coffee mug on my developer desk",
                "Cynical",
                trace_dir=Path(tmp_dir),
            )
            saved_path = Path(result.trace_path)

            self.assertTrue(saved_path.exists())
            self.assertEqual(result.object_understanding.object.name, "coffee mug")
            self.assertEqual(saved_path.stem, result.trace.trace_id)

    def test_pipeline_records_minicpm_vision_runtime(self) -> None:
        response = """
        {"object":{"name":"desk lamp","visible_features":["metal shade","thin neck","warm light"],"likely_context":"desk","confidence":0.91}}
        """
        env = {
            "OBJECTVERSE_VISION_BACKEND": "minicpm-v",
            "VISION_MODEL_ID": "openbmb/MiniCPM-V-2_6",
            "OBJECTVERSE_TEXT_BACKEND": "mock",
        }

        with (
            patch.dict("os.environ", env, clear=False),
            patch("src.models.vision_runner._load_rgb_image", return_value=object()),
            patch("src.models.vision_runner._load_minicpm_components", return_value=(FakeMiniCpmModel(response), object())),
        ):
            result = generate_object_diary("/tmp/lamp.png", "desk lamp", "Dramatic", save=False)

        self.assertEqual(result.object_understanding.object.name, "desk lamp")
        self.assertEqual(result.trace.model_runtime["vision"], "minicpm-v object understanding")
        self.assertIn("mock-text-runtime", result.trace.fallbacks)
        self.assertNotIn("mock-runtime", result.trace.fallbacks)

    def test_chat_uses_current_persona(self) -> None:
        object_understanding = understand_object(None, "dusty black mechanical keyboard")
        persona = generate_persona(object_understanding, "Philosopher")
        reply = reply_as_object(persona.model_dump(mode="json"), "What have you seen?")

        self.assertIn(persona.persona.character_name, reply)
        self.assertIn("keyboard", reply)

    def test_trace_save_and_anonymization(self) -> None:
        description = "old phone from user@example.com with serial 123456789"
        object_understanding = understand_object(None, description)
        persona = generate_persona(object_understanding, "Lonely")
        diary = generate_diary(persona, "Lonely")
        trace = build_trace(None, description, "Lonely", object_understanding, persona, diary)

        self.assertIn("[redacted-email]", anonymize_text(description))
        self.assertIn("[redacted-number]", trace.input["description"])
        self.assertIn("mock-runtime", trace.fallbacks)

        with tempfile.TemporaryDirectory() as tmp_dir:
            path = Path(save_trace(trace, Path(tmp_dir)))
            saved = json.loads(path.read_text(encoding="utf-8"))

        self.assertEqual(saved["trace_id"], trace.trace_id)
        self.assertEqual(saved["model_runtime"]["vision"], "mock object understanding")

    def test_generate_sample_traces_script(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            paths = generate_sample_traces(Path(tmp_dir))
            saved = [json.loads(path.read_text(encoding="utf-8")) for path in paths]

        self.assertEqual(len(paths), 6)
        self.assertEqual(saved[0]["trace_id"], "sample-01")
        self.assertTrue(all(item["fallbacks"] == ["mock-runtime"] for item in saved))

    def test_initial_stage_acceptance_checks(self) -> None:
        results = run_checks()

        self.assertIn("required files exist", results)
        self.assertIn("six sample traces validate", results)
        self.assertIn("Gradio Blocks app builds", results)


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