207 lines
6.8 KiB
Python
207 lines
6.8 KiB
Python
"""Tests for editing archived scheduled-task run files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from deerflow.config import paths as paths_module
|
|
from deerflow.persistence.scheduled_tasks import MemoryScheduledTaskStore
|
|
from deerflow.runtime.scheduler.service import ScheduledTaskService
|
|
|
|
|
|
@pytest.fixture
|
|
def isolated_paths(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
|
|
monkeypatch.setattr(paths_module, "_paths", None)
|
|
yield tmp_path
|
|
monkeypatch.setattr(paths_module, "_paths", None)
|
|
|
|
|
|
@pytest.fixture
|
|
def store() -> MemoryScheduledTaskStore:
|
|
return MemoryScheduledTaskStore()
|
|
|
|
|
|
def _build_service(store: MemoryScheduledTaskStore) -> ScheduledTaskService:
|
|
service = ScheduledTaskService(app=None, store=store, poll_interval_seconds=60.0)
|
|
service._element_client = None # noqa: SLF001
|
|
return service
|
|
|
|
|
|
def _seed_task(store: MemoryScheduledTaskStore, *, task_id: str, user_id: str, published: bool = False) -> None:
|
|
store._tasks[task_id] = { # noqa: SLF001
|
|
"task_id": task_id,
|
|
"user_id": user_id,
|
|
"scheduler_thread_id": f"thread-{user_id}",
|
|
"name": "demo",
|
|
"prompt": "hi",
|
|
"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,
|
|
}
|
|
|
|
|
|
async def _seed_run_with_file(
|
|
store: MemoryScheduledTaskStore,
|
|
isolated_paths: Path,
|
|
*,
|
|
run_id: str,
|
|
task_id: str,
|
|
user_id: str,
|
|
file_name: str,
|
|
file_text: str,
|
|
mime_type: str = "text/markdown",
|
|
) -> dict[str, Any]:
|
|
archive_dir = isolated_paths / "scheduled-task-runs" / run_id / "files"
|
|
archive_dir.mkdir(parents=True)
|
|
file_path = archive_dir / file_name
|
|
file_path.write_text(file_text, encoding="utf-8")
|
|
|
|
file_entry = {
|
|
"name": file_name,
|
|
"path": str(file_path),
|
|
"mime_type": mime_type,
|
|
"kind": "markdown" if file_name.endswith(".md") else "artifact",
|
|
"size": file_path.stat().st_size,
|
|
}
|
|
result = {"title": "demo", "content": "ignored", "format": "markdown", "files": [file_entry]}
|
|
|
|
scheduled_for = datetime.now(UTC)
|
|
await store.create_run_once(
|
|
{
|
|
"id": run_id,
|
|
"task_id": task_id,
|
|
"user_id": user_id,
|
|
"scheduler_thread_id": f"thread-{user_id}",
|
|
"scheduled_for": scheduled_for,
|
|
"status": "succeeded",
|
|
"started_at": scheduled_for,
|
|
}
|
|
)
|
|
await store.update_run(run_id, {"result": result, "finished_at": scheduled_for})
|
|
return file_entry
|
|
|
|
|
|
# ---------- can_edit flag on get_run ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_run_marks_owner_as_can_edit(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice", published=True)
|
|
await _seed_run_with_file(store, isolated_paths, run_id="r1", task_id="t1", user_id="alice", file_name="a.md", file_text="# hi")
|
|
service = _build_service(store)
|
|
|
|
run = await service.get_run("alice", "r1")
|
|
assert run is not None
|
|
assert run["can_edit"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_run_marks_subscriber_as_read_only(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice", published=True)
|
|
await _seed_run_with_file(store, isolated_paths, run_id="r1", task_id="t1", user_id="alice", file_name="a.md", file_text="# hi")
|
|
service = _build_service(store)
|
|
await service.subscribe_task("bob", "t1", notify_element=False)
|
|
|
|
run = await service.get_run("bob", "r1")
|
|
assert run is not None
|
|
assert run["can_edit"] is False
|
|
|
|
|
|
# ---------- update_run_file ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_owner_can_overwrite_md_file(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice")
|
|
file_entry = await _seed_run_with_file(
|
|
store,
|
|
isolated_paths,
|
|
run_id="r1",
|
|
task_id="t1",
|
|
user_id="alice",
|
|
file_name="report.md",
|
|
file_text="# original",
|
|
)
|
|
service = _build_service(store)
|
|
|
|
new_text = "# updated\n\nA new paragraph."
|
|
meta = await service.update_run_file("alice", "r1", 0, new_text)
|
|
|
|
assert meta["name"] == "report.md"
|
|
assert meta["size"] == len(new_text.encode("utf-8"))
|
|
|
|
# The on-disk file is overwritten.
|
|
assert Path(file_entry["path"]).read_text(encoding="utf-8") == new_text
|
|
|
|
# The size in the run row is refreshed.
|
|
run = await service.get_run("alice", "r1")
|
|
assert run is not None
|
|
assert run["result"]["files"][0]["size"] == meta["size"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subscriber_cannot_edit(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice", published=True)
|
|
await _seed_run_with_file(store, isolated_paths, run_id="r1", task_id="t1", user_id="alice", file_name="a.md", file_text="# hi")
|
|
service = _build_service(store)
|
|
await service.subscribe_task("bob", "t1", notify_element=False)
|
|
|
|
with pytest.raises(PermissionError):
|
|
await service.update_run_file("bob", "r1", 0, "# evil")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unrelated_user_gets_not_found(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice")
|
|
await _seed_run_with_file(store, isolated_paths, run_id="r1", task_id="t1", user_id="alice", file_name="a.md", file_text="# hi")
|
|
service = _build_service(store)
|
|
|
|
with pytest.raises(PermissionError):
|
|
await service.update_run_file("eve", "r1", 0, "# evil")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_run_returns_filenotfound(store, isolated_paths):
|
|
service = _build_service(store)
|
|
with pytest.raises(FileNotFoundError):
|
|
await service.update_run_file("alice", "missing-run", 0, "# x")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_file_index_rejected(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice")
|
|
await _seed_run_with_file(store, isolated_paths, run_id="r1", task_id="t1", user_id="alice", file_name="a.md", file_text="# hi")
|
|
service = _build_service(store)
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
await service.update_run_file("alice", "r1", 5, "# x")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_text_file_rejected(store, isolated_paths):
|
|
_seed_task(store, task_id="t1", user_id="alice")
|
|
await _seed_run_with_file(
|
|
store,
|
|
isolated_paths,
|
|
run_id="r1",
|
|
task_id="t1",
|
|
user_id="alice",
|
|
file_name="diagram.png",
|
|
file_text="not really a png but stored as such",
|
|
mime_type="image/png",
|
|
)
|
|
service = _build_service(store)
|
|
|
|
with pytest.raises(ValueError, match="text-previewable"):
|
|
await service.update_run_file("alice", "r1", 0, "# x")
|