"""Tests for scheduled task publish/subscribe + mirrored Element delivery.""" from __future__ import annotations from datetime import UTC, datetime from typing import Any import pytest from deerflow.persistence.scheduled_tasks import MemoryScheduledTaskStore from deerflow.runtime.scheduler.service import ScheduledTaskService @pytest.fixture def store() -> MemoryScheduledTaskStore: return MemoryScheduledTaskStore() def _seed_task( store: MemoryScheduledTaskStore, *, task_id: str, user_id: str, name: str = "demo", prompt: str = "Hello world", published: bool = False, ) -> dict[str, Any]: task = { "task_id": task_id, "user_id": user_id, "scheduler_thread_id": f"thread-{user_id}", "name": name, "prompt": prompt, "schedule_text": "0 9 * * *", "cron_expr": "0 9 * * *", "timezone": "UTC", "enabled": True, "next_run_at": datetime.now(UTC), "execution_context": {}, "published": published, "published_at": datetime.now(UTC) if published else None, } store._tasks[task_id] = task # noqa: SLF001 return task def _build_service(store: MemoryScheduledTaskStore) -> ScheduledTaskService: # Build with no FastAPI app — the methods we exercise here don't touch app.state. service = ScheduledTaskService(app=None, store=store, poll_interval_seconds=60.0) service._element_client = None # noqa: SLF001 — disable real delivery return service @pytest.mark.asyncio async def test_publish_unpublish_only_owner(store): task = _seed_task(store, task_id="t1", user_id="alice") service = _build_service(store) assert await service.publish_task("bob", task["task_id"]) is None, "non-owner cannot publish" published = await service.publish_task("alice", task["task_id"]) assert published is not None assert published["published"] is True assert published["published_at"] is not None listed = await service.list_published(limit=10, offset=0) assert len(listed) == 1 assert listed[0]["task_id"] == "t1" assert listed[0]["prompt_summary"] == "Hello world" assert "prompt" not in listed[0], "raw prompt must not leak through summary endpoint" unpublished = await service.unpublish_task("alice", task["task_id"]) assert unpublished["published"] is False assert (await service.list_published()) == [] @pytest.mark.asyncio async def test_subscribe_requires_published_for_first_time(store): task = _seed_task(store, task_id="t1", user_id="alice", published=False) service = _build_service(store) # Cannot subscribe to a private task. assert await service.subscribe_task("bob", task["task_id"], notify_element=False) is None # Author publishes; bob subscribes. await service.publish_task("alice", task["task_id"]) sub = await service.subscribe_task("bob", task["task_id"], notify_element=True) assert sub is not None and sub["notify_element"] is True # Author cannot subscribe to their own task. assert await service.subscribe_task("alice", task["task_id"], notify_element=True) is None @pytest.mark.asyncio async def test_unpublish_keeps_existing_subscribers(store): task = _seed_task(store, task_id="t1", user_id="alice", published=True) service = _build_service(store) sub = await service.subscribe_task("bob", task["task_id"], notify_element=False) assert sub is not None # Author un-publishes — bob keeps acquired access. await service.unpublish_task("alice", task["task_id"]) assert await store.get_subscription(task["task_id"], "bob") is not None # New user cannot subscribe after un-publish. assert await service.subscribe_task("carol", task["task_id"], notify_element=False) is None # Existing subscriber can still toggle their notification preference. updated = await service.subscribe_task("bob", task["task_id"], notify_element=True) assert updated is not None and updated["notify_element"] is True @pytest.mark.asyncio async def test_list_tasks_includes_subscribed_with_role(store): _seed_task(store, task_id="t1", user_id="alice", published=True) _seed_task(store, task_id="t2", user_id="bob") service = _build_service(store) await service.subscribe_task("bob", "t1", notify_element=False) bobs = await service.list_tasks("bob") by_id = {t["task_id"]: t for t in bobs} assert by_id["t2"]["role"] == "owner" assert by_id["t1"]["role"] == "subscriber" assert by_id["t1"]["_subscription"]["notify_element"] is False @pytest.mark.asyncio async def test_subscriber_can_read_runs_and_get_run(store): _seed_task(store, task_id="t1", user_id="alice", published=True) service = _build_service(store) await service.subscribe_task("bob", "t1", notify_element=False) # Seed a run owned by alice. scheduled_for = datetime.now(UTC) run = await store.create_run_once( { "id": "run-1", "task_id": "t1", "user_id": "alice", "scheduler_thread_id": "thread-alice", "scheduled_for": scheduled_for, "status": "succeeded", "started_at": scheduled_for, } ) assert run is not None alice_runs = await service.list_runs("alice", "t1", limit=10) assert alice_runs[0]["can_edit"] is True # Bob (subscriber) sees runs on the task. bob_runs = await service.list_runs("bob", "t1", limit=10) assert [r["id"] for r in bob_runs] == ["run-1"] assert bob_runs[0]["can_edit"] is False # Bob can fetch the run by id. fetched = await service.get_run("bob", "run-1") assert fetched is not None and fetched["id"] == "run-1" assert fetched["can_edit"] is False # Bob can read subscribed records, but cannot delete Alice's record. with pytest.raises(PermissionError): await service.delete_run("bob", "run-1") assert await store.get_run("run-1", "alice") is not None # Carol (no subscription) cannot. assert await service.list_runs("carol", "t1", limit=10) == [] assert await service.get_run("carol", "run-1") is None @pytest.mark.asyncio async def test_unsubscribe_removes_visibility(store): _seed_task(store, task_id="t1", user_id="alice", published=True) service = _build_service(store) await service.subscribe_task("bob", "t1", notify_element=False) assert await service.unsubscribe_task("bob", "t1") is True bob_tasks = await service.list_tasks("bob") assert all(t["task_id"] != "t1" for t in bob_tasks) @pytest.mark.asyncio async def test_mirror_element_delivery_to_subscribers(store): _seed_task(store, task_id="t1", user_id="alice", published=True) service = _build_service(store) # Author and two subscribers configure delivery profiles. await store.update_delivery_profile("alice", {"element_user_id": "@alice:example"}) await store.update_delivery_profile("bob", {"element_user_id": "@bob:example"}) await store.update_delivery_profile("carol", {"element_user_id": "@carol:example"}) await service.subscribe_task("bob", "t1", notify_element=True) await service.subscribe_task("carol", "t1", notify_element=False) # opted out sent: list[tuple[str, str]] = [] class _FakeElement: async def send_direct_message(self, user_id: str, body: str, room_id: str | None = None) -> str: sent.append((user_id, body)) return f"!room-{user_id}" service._element_client = _FakeElement() # noqa: SLF001 task = await store.get_task_any("t1") assert task is not None await service._deliver_element_result(task, "result body") # noqa: SLF001 recipients = sorted(user for user, _ in sent) assert recipients == ["@alice:example", "@bob:example"], "only author + opted-in subscriber receive" # All recipients get the same body content (we sent identical messages). bodies = {body for _, body in sent} assert bodies == {"【定时任务】demo\n\nresult body"} # Carol is not in the recipient list. assert all(uid != "@carol:example" for uid, _ in sent) @pytest.mark.asyncio async def test_published_summary_truncates_long_prompt(store): long_prompt = "abcdef" * 100 # 600 chars _seed_task(store, task_id="t1", user_id="alice", prompt=long_prompt, published=True) service = _build_service(store) listed = await service.list_published() summary = listed[0]["prompt_summary"] assert len(summary) <= 201 # 200 chars + ellipsis assert summary.endswith("…")