jiangchang-platform-kit/tests/test_activity_stream.py

# -*- coding: utf-8 -*-
"""SAS stream_start / stream_delta / stream_end / stream_channel."""
from __future__ import annotations

import json
import os
import shutil
import tempfile
import unittest
from pathlib import Path

from jiangchang_skill_core import activity
from jiangchang_skill_core.runtime_env import get_runs_dir


class _ActivityEnvIsolation(unittest.TestCase):
    _KEYS = (
        "JIANGCHANG_DATA_ROOT",
        "JIANGCHANG_USER_ID",
        "JIANGCHANG_JOB_ID",
        "JIANGCHANG_RUN_MODE",
        "JIANGCHANG_ACTIVITY_DEBUG",
        "JIANGCHANG_STREAM_COALESCE_MS",
        "JIANGCHANG_STREAM_COALESCE_CHARS",
        "JIANGCHANG_CONTROL_DISABLED",
    )

    def setUp(self) -> None:
        self._saved = {k: os.environ.get(k) for k in self._KEYS}
        self._tmps: list[str] = []
        for k in self._KEYS:
            os.environ.pop(k, None)
        os.environ["JIANGCHANG_CONTROL_DISABLED"] = "1"
        activity._reset_for_tests()

    def tearDown(self) -> None:
        activity._reset_for_tests()
        for k, v in self._saved.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v
        for tmp in self._tmps:
            shutil.rmtree(tmp, ignore_errors=True)

    def _tmp_data_root(self) -> str:
        tmp = tempfile.mkdtemp()
        self._tmps.append(tmp)
        return tmp

    def _events(self, job_id: str) -> list[dict]:
        journal = Path(get_runs_dir()) / f"{job_id}.jsonl"
        lines = journal.read_text(encoding="utf-8").strip().splitlines()
        return [json.loads(line) for line in lines]


class TestActivityStream(_ActivityEnvIsolation):
    def test_stream_writes_start_delta_end_without_step_gate(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "job-stream-1"
        os.environ["JIANGCHANG_STREAM_COALESCE_MS"] = "0"
        activity._reset_for_tests()

        activity.stream_start("answer", skill="demo")
        activity.stream_delta("你", stream_id="answer", skill="demo")
        activity.stream_delta("好", stream_id="answer", skill="demo")
        activity.stream_end("answer", skill="demo")

        events = self._events("job-stream-1")
        stream_events = [e for e in events if e.get("type") == "stream"]
        phases = [e.get("phase") for e in stream_events]
        self.assertEqual(phases[0], "start")
        self.assertEqual(phases[-1], "end")
        deltas = [e.get("delta") for e in stream_events if e.get("phase") == "delta"]
        self.assertEqual("".join(deltas), "你好")
        self.assertEqual(activity._step_index, 0)

    def test_stream_channel_context_manager(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "job-stream-ch"
        os.environ["JIANGCHANG_STREAM_COALESCE_MS"] = "0"
        activity._reset_for_tests()

        with activity.stream_channel("answer", skill="demo") as ch:
            ch.delta("A")
            ch.delta("B", flush=True)

        events = self._events("job-stream-ch")
        stream_events = [e for e in events if e.get("type") == "stream"]
        self.assertEqual(stream_events[0]["phase"], "start")
        self.assertEqual(stream_events[-1]["phase"], "end")
        deltas = [e.get("delta") for e in stream_events if e.get("phase") == "delta"]
        self.assertEqual("".join(d for d in deltas if d), "AB")

    def test_stream_delta_auto_starts(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "job-stream-auto"
        os.environ["JIANGCHANG_STREAM_COALESCE_MS"] = "0"
        activity._reset_for_tests()

        activity.stream_delta("x", stream_id="answer")
        activity.stream_end("answer")

        events = self._events("job-stream-auto")
        stream_events = [e for e in events if e.get("type") == "stream"]
        self.assertEqual(stream_events[0]["phase"], "start")
        self.assertTrue(any(e.get("phase") == "delta" and e.get("delta") == "x" for e in stream_events))

    def test_empty_delta_ignored(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "job-stream-empty"
        os.environ["JIANGCHANG_STREAM_COALESCE_MS"] = "0"
        activity._reset_for_tests()

        activity.stream_start("answer")
        activity.stream_delta("", stream_id="answer")
        activity.stream_end("answer")

        events = self._events("job-stream-empty")
        stream_events = [e for e in events if e.get("type") == "stream"]
        self.assertFalse(any(e.get("phase") == "delta" for e in stream_events))

    def test_public_exports(self) -> None:
        from jiangchang_skill_core import (
            StreamChannel,
            stream_channel,
            stream_delta,
            stream_end,
            stream_start,
        )

        self.assertTrue(callable(stream_start))
        self.assertTrue(callable(stream_delta))
        self.assertTrue(callable(stream_end))
        self.assertTrue(callable(stream_channel))
        self.assertTrue(issubclass(StreamChannel, object))


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