tests/test_observability_and_progress.py

import os
import tempfile
import unittest

import core
import model_client
import story_store
from observability import RunProfiler, resource_snapshot
from schemas import GraftPatch, GraftPlan

from test_app_flow import graft_patch, graft_plan, initial_payload


class ResolveExecutionTests(unittest.TestCase):
    def setUp(self):
        self._saved = {k: os.environ.get(k) for k in ("BQ_DEVICE", "SPACE_ID")}

    def tearDown(self):
        for key, value in self._saved.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value

    def test_explicit_modes_are_honored(self):
        for mode in ("cpu", "mps", "cuda", "zerogpu"):
            os.environ["BQ_DEVICE"] = mode
            self.assertEqual(model_client.resolve_execution(), mode)

    def test_unknown_mode_falls_back_without_crashing(self):
        os.environ["BQ_DEVICE"] = "banana"
        os.environ.pop("SPACE_ID", None)
        # auto detection: never raises, always returns a concrete mode.
        self.assertIn(model_client.resolve_execution(), {"cuda", "mps", "cpu", "zerogpu"})

    def test_auto_on_space_is_zerogpu(self):
        os.environ["BQ_DEVICE"] = "auto"
        os.environ["SPACE_ID"] = "user/space"
        # _is_zerogpu only fires when the `spaces` runtime is importable.
        expected = "zerogpu" if model_client.spaces is not None else None
        if expected:
            self.assertEqual(model_client.resolve_execution(), "zerogpu")


class ResourceSnapshotTests(unittest.TestCase):
    def test_snapshot_returns_dict_and_never_raises(self):
        snapshot = resource_snapshot()
        self.assertIsInstance(snapshot, dict)
        # psutil is a dependency, so process memory should always be present.
        self.assertIn("rss_mb", snapshot)

    def test_profiler_summary_runs(self):
        profiler = RunProfiler("test", label="x=1")
        with profiler.stage("a"):
            pass
        profiler.note_message()
        self.assertEqual(profiler.messages, 1)
        profiler.summary()  # must not raise


class StitchProgressTests(unittest.TestCase):
    def test_on_progress_streams_monotonic_events_and_returns_result(self):
        saved_data_dir = os.environ.get("DATA_DIR")
        saved_store_gen = story_store.generate_json
        saved_core_gen = core.generate_json
        saved_cache = model_client._execution_cache
        try:
            with tempfile.TemporaryDirectory() as directory:
                os.environ["DATA_DIR"] = directory
                # Force the local (non-zerogpu) path so on_step is plumbed through.
                model_client._execution_cache = "cpu"
                story_store.generate_json = lambda *a, **k: initial_payload()

                story = core.create("A lighthouse seed.")

                def fake_generate(*args, **kwargs):
                    schema_model = args[1]
                    on_step = kwargs.get("on_step")
                    if on_step:  # exercise the token-progress callback
                        total = kwargs.get("max_new_tokens", 100)
                        on_step(total // 2, total)
                    if schema_model is GraftPlan:
                        return graft_plan()
                    if schema_model is GraftPatch:
                        return graft_patch()
                    raise AssertionError(schema_model)

                core.generate_json = fake_generate

                events = []
                result = core.stitch(
                    story.story_id,
                    "Blue glass reflects one day late.",
                    on_progress=events.append,
                )

                # The synchronous contract is preserved.
                self.assertEqual(result.story.graft_count, 1)

                progress = [e for e in events if e["type"] == "progress"]
                self.assertTrue(progress)
                # Fractions never go backwards and finish at 1.0.
                fractions = [e["fraction"] for e in progress]
                self.assertEqual(fractions, sorted(fractions))
                self.assertEqual(progress[-1]["fraction"], 1.0)
                # Both stages were observed and counted as messages.
                stages = {e["stage"] for e in progress}
                self.assertIn("planning", stages)
                self.assertIn("writing", stages)
                self.assertEqual(progress[-1]["messagesProcessed"], 2)
                # Token-level events carry token counts.
                token_events = [e for e in progress if e["tokensTotal"]]
                self.assertTrue(token_events)
        finally:
            story_store.generate_json = saved_store_gen
            core.generate_json = saved_core_gen
            model_client._execution_cache = saved_cache
            if saved_data_dir is None:
                os.environ.pop("DATA_DIR", None)
            else:
                os.environ["DATA_DIR"] = saved_data_dir


class ZeroGpuFallbackTests(unittest.TestCase):
    def test_quota_error_detection(self):
        import gradio as gr

        import app

        quota = gr.Error("You have exceeded your free ZeroGPU quota (300s requested vs. 0s left).")
        credits = gr.Error("ZeroGPU pending credits exceeded.")
        self.assertTrue(app._is_quota_error(quota))
        self.assertTrue(app._is_quota_error(credits))
        self.assertFalse(app._is_quota_error(gr.Error("The bindery hit an internal error.")))
        self.assertFalse(app._is_quota_error(ValueError("not a gradio error")))

    def test_zerogpu_quota_miss_falls_back_to_cpu_streaming(self):
        import gradio as gr

        import app

        saved_mode = model_client._execution_cache
        saved_stitch = core.stitch
        saved_result_event = app._result_event
        try:
            model_client._execution_cache = "zerogpu"
            forced = []

            def fake_stitch(story_id, fragment, on_progress=None, force_cpu=False):
                forced.append(force_cpu)
                if not force_cpu:
                    # The per-user GPU quota is spent.
                    raise gr.Error("You have exceeded your ZeroGPU quota. Try again in 1:00:00.")
                if on_progress:
                    on_progress(
                        {
                            "type": "progress",
                            "stage": "planning",
                            "phase": "reading",
                            "label": "Reading",
                            "stageIndex": 1,
                            "stageTotal": 2,
                            "fraction": 0.1,
                            "tokensDone": 400,
                            "tokensTotal": 4096,
                            "etaSeconds": 12.0,
                            "messagesProcessed": 0,
                        }
                    )
                return object()  # stand-in AppliedPatchResult

            core.stitch = fake_stitch
            app._result_event = lambda result: {"type": "result", "ok": True}

            events = list(app._stitch_events("abc123", "a fragment"))

            # GPU attempt first (force_cpu False), then CPU retry (force_cpu True).
            self.assertEqual(forced, [False, True])
            progress = [e for e in events if e["type"] == "progress"]
            self.assertTrue(progress)
            self.assertEqual(progress[0]["notice"], app._CPU_FALLBACK_NOTICE)
            self.assertEqual(events[-1], {"type": "result", "ok": True})
        finally:
            model_client._execution_cache = saved_mode
            core.stitch = saved_stitch
            app._result_event = saved_result_event


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