deerflow-code/offline-backend-20260512/backend/tests/test_scheduled_tasks_publish.py

"""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("…")