jiangchang-platform-kit/tests/test_activity_control.py

# -*- coding: utf-8 -*-
"""SRCP v1: control.json, checkpoint, pause/resume/stop, interruptible_sleep."""
from __future__ import annotations

import asyncio
import io
import json
import os
import shutil
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from jiangchang_skill_core import activity
from jiangchang_skill_core.activity import JobStopped
from jiangchang_skill_core.runtime_env import get_runs_dir


class _ControlEnvIsolation(unittest.TestCase):
    _KEYS = (
        "JIANGCHANG_DATA_ROOT",
        "JIANGCHANG_USER_ID",
        "JIANGCHANG_JOB_ID",
        "JIANGCHANG_RUN_MODE",
        "JIANGCHANG_CONTROL_DISABLED",
        "JIANGCHANG_CONTROL_POLL_SEC",
    )

    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)
        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 _write_control(self, job_id: str, command: str, *, nonce: int = 1) -> None:
        path = Path(get_runs_dir()) / f"{job_id}.control.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(
            json.dumps(
                {
                    "schema_version": 1,
                    "command": command,
                    "reason": "user",
                    "ts": 1_700_000_000_000,
                    "nonce": nonce,
                },
                ensure_ascii=False,
            ),
            encoding="utf-8",
        )

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


class TestControlProtocol(_ControlEnvIsolation):
    def test_get_control_path(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-path"
        activity._reset_for_tests()
        self.assertTrue(activity.get_control_path().endswith("ctrl-path.control.json"))

    def test_stop_raises_job_stopped(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-stop"
        os.environ["JIANGCHANG_CONTROL_POLL_SEC"] = "0.05"
        activity._reset_for_tests()
        self._write_control("ctrl-stop", "stop")

        with self.assertRaises(JobStopped):
            activity.checkpoint()

        events = self._journal_events("ctrl-stop")
        lifecycle = [e for e in events if e.get("type") == "lifecycle"]
        self.assertTrue(any(e.get("status") == "stopping" for e in lifecycle))

    def test_emit_stop_before_activity_write(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-emit-stop"
        activity._reset_for_tests()
        self._write_control("ctrl-emit-stop", "stop")

        with self.assertRaises(JobStopped):
            activity.emit("不应写入")

        events = self._journal_events("ctrl-emit-stop")
        self.assertFalse(any(e.get("type") == "activity" for e in events))
        self.assertTrue(any(e.get("type") == "lifecycle" and e.get("status") == "stopping" for e in events))

    def test_pause_resume_via_emit(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-pause"
        os.environ["JIANGCHANG_CONTROL_POLL_SEC"] = "0.05"
        activity._reset_for_tests()
        self._write_control("ctrl-pause", "pause")

        import threading
        import time as time_mod

        def resume_later() -> None:
            time_mod.sleep(0.15)
            self._write_control("ctrl-pause", "resume", nonce=2)

        threading.Thread(target=resume_later, daemon=True).start()
        activity.emit("步骤一", skill="demo")

        events = self._journal_events("ctrl-pause")
        statuses = [e.get("status") for e in events if e.get("type") == "lifecycle"]
        self.assertIn("paused", statuses)
        self.assertIn("resumed", statuses)
        activity_events = [e for e in events if e.get("type") == "activity"]
        self.assertEqual(activity_events[-1]["text"], "步骤一")
        lifecycle = [e for e in events if e.get("type") == "lifecycle"]
        self.assertEqual(lifecycle[0].get("step_index"), 1)

    def test_rpa_step_gate_increments_step_index(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-rpa"
        activity._reset_for_tests()

        @activity.rpa_step("打开页面")
        def open_page() -> None:
            return None

        open_page()
        events = self._journal_events("ctrl-rpa")
        lifecycle_or_activity = [e for e in events if e.get("type") in ("activity", "lifecycle")]
        self.assertGreaterEqual(len(lifecycle_or_activity), 2)

    def test_job_context_finishes_on_job_stopped(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        activity._reset_for_tests()

        with patch("sys.stdout", io.StringIO()):
            with self.assertRaises(JobStopped):
                with activity.job_context(job_id="ctx-stop", skill="demo-skill"):
                    self._write_control("ctx-stop", "stop")
                    activity.checkpoint()

        journal = Path(get_runs_dir()) / "ctx-stop.jsonl"
        lines = journal.read_text(encoding="utf-8").strip().splitlines()
        last = json.loads(lines[-1])
        self.assertEqual(last["type"], "error")
        self.assertEqual(last.get("error_code"), "JOB_STOPPED")

    def test_interruptible_sleep_honors_stop(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-sleep"
        os.environ["JIANGCHANG_CONTROL_POLL_SEC"] = "0.05"
        activity._reset_for_tests()
        self._write_control("ctrl-sleep", "stop")

        async def run() -> None:
            with self.assertRaises(JobStopped):
                await activity.interruptible_sleep(5.0)

        asyncio.run(run())

    def test_interruptible_sleep_sync_honors_pause(self) -> None:
        tmp = self._tmp_data_root()
        os.environ["JIANGCHANG_DATA_ROOT"] = tmp
        os.environ["JIANGCHANG_JOB_ID"] = "ctrl-sleep-sync"
        os.environ["JIANGCHANG_CONTROL_POLL_SEC"] = "0.05"
        activity._reset_for_tests()
        self._write_control("ctrl-sleep-sync", "pause")

        def resume_later() -> None:
            import threading
            import time

            def _resume() -> None:
                time.sleep(0.15)
                self._write_control("ctrl-sleep-sync", "resume", nonce=2)

            threading.Thread(target=_resume, daemon=True).start()

        resume_later()
        activity.interruptible_sleep_sync(0.4, chunk=0.05)

        events = self._journal_events("ctrl-sleep-sync")
        statuses = [e.get("status") for e in events if e.get("type") == "lifecycle"]
        self.assertIn("paused", statuses)
        self.assertIn("resumed", statuses)


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