mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-08 22:22:17 +02:00
test(chat): cover watch tools
This commit is contained in:
parent
f1ea60365c
commit
f0c8d9ba5c
4 changed files with 419 additions and 0 deletions
|
|
@ -0,0 +1,182 @@
|
|||
"""Integration tests for the watch chat tools against real service + Postgres.
|
||||
|
||||
The tools are thin adapters the intelligence agent calls; here they run their
|
||||
real path — open a session, build ``AutomationService``, and create / find /
|
||||
delete / enqueue — against real persistence (RBAC + model gate included), with
|
||||
only the Celery enqueue spied. This is what proves ``start_watch`` actually
|
||||
binds a watch, ``stop_watch`` removes it, and ``refresh_watch`` enqueues a run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools.refresh_watch import (
|
||||
create_refresh_watch_tool,
|
||||
)
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools.start_watch import (
|
||||
create_start_watch_tool,
|
||||
)
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools.stop_watch import (
|
||||
create_stop_watch_tool,
|
||||
)
|
||||
from app.auth.context import AuthContext
|
||||
from app.automations.services.automation import AutomationService
|
||||
from app.automations.services.chat_watch import find_watches_for_thread
|
||||
from app.db import ChatVisibility, NewChatThread, User, Workspace
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
_CRON = "0 9 * * 1-5"
|
||||
_TZ = "UTC"
|
||||
|
||||
|
||||
def _billable(workspace: Workspace) -> None:
|
||||
workspace.chat_model_id = 1
|
||||
workspace.image_gen_model_id = 1
|
||||
workspace.vision_model_id = 1
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def thread(
|
||||
db_session: AsyncSession, db_user: User, db_workspace: Workspace
|
||||
) -> NewChatThread:
|
||||
row = NewChatThread(
|
||||
title="Watched chat",
|
||||
workspace_id=db_workspace.id,
|
||||
created_by_id=db_user.id,
|
||||
visibility=ChatVisibility.PRIVATE,
|
||||
)
|
||||
db_session.add(row)
|
||||
await db_session.flush()
|
||||
return row
|
||||
|
||||
|
||||
async def test_start_watch_tool_binds_a_watch_to_the_chat(
|
||||
tools_use_test_session,
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
|
||||
tool = create_start_watch_tool(
|
||||
workspace_id=db_workspace.id,
|
||||
thread_id=thread.id,
|
||||
auth_context=AuthContext.session(db_user),
|
||||
)
|
||||
out = await tool.ainvoke(
|
||||
{"message": "what changed?", "cron": _CRON, "timezone": _TZ}
|
||||
)
|
||||
|
||||
assert out["status"] == "watching"
|
||||
assert isinstance(out["automation_id"], int)
|
||||
|
||||
service = AutomationService(session=db_session, auth=AuthContext.session(db_user))
|
||||
found = await find_watches_for_thread(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread.id
|
||||
)
|
||||
assert [a.id for a in found] == [out["automation_id"]]
|
||||
|
||||
|
||||
async def test_start_watch_tool_reports_error_when_models_not_billable(
|
||||
tools_use_test_session,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
# Default workspace models are Auto/None -> the automation gate rejects it.
|
||||
tool = create_start_watch_tool(
|
||||
workspace_id=db_workspace.id,
|
||||
thread_id=thread.id,
|
||||
auth_context=AuthContext.session(db_user),
|
||||
)
|
||||
out = await tool.ainvoke(
|
||||
{"message": "track it", "cron": _CRON, "timezone": _TZ}
|
||||
)
|
||||
assert out["status"] == "error"
|
||||
|
||||
|
||||
async def test_stop_watch_tool_removes_the_chats_watch(
|
||||
tools_use_test_session,
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
auth = AuthContext.session(db_user)
|
||||
|
||||
start = create_start_watch_tool(
|
||||
workspace_id=db_workspace.id, thread_id=thread.id, auth_context=auth
|
||||
)
|
||||
await start.ainvoke({"message": "m", "cron": _CRON, "timezone": _TZ})
|
||||
|
||||
stop = create_stop_watch_tool(
|
||||
workspace_id=db_workspace.id, thread_id=thread.id, auth_context=auth
|
||||
)
|
||||
out = await stop.ainvoke({})
|
||||
|
||||
assert out["status"] == "stopped"
|
||||
assert out["count"] == 1
|
||||
|
||||
service = AutomationService(session=db_session, auth=auth)
|
||||
found = await find_watches_for_thread(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread.id
|
||||
)
|
||||
assert found == []
|
||||
|
||||
|
||||
async def test_stop_watch_tool_reports_not_watching_when_none(
|
||||
tools_use_test_session,
|
||||
db_workspace: Workspace,
|
||||
db_user: User,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
stop = create_stop_watch_tool(
|
||||
workspace_id=db_workspace.id,
|
||||
thread_id=thread.id,
|
||||
auth_context=AuthContext.session(db_user),
|
||||
)
|
||||
out = await stop.ainvoke({})
|
||||
assert out["status"] == "not_watching"
|
||||
|
||||
|
||||
async def test_refresh_watch_tool_enqueues_a_run(
|
||||
tools_use_test_session,
|
||||
enqueue_spy: list[dict],
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
auth = AuthContext.session(db_user)
|
||||
|
||||
start = create_start_watch_tool(
|
||||
workspace_id=db_workspace.id, thread_id=thread.id, auth_context=auth
|
||||
)
|
||||
started = await start.ainvoke({"message": "m", "cron": _CRON, "timezone": _TZ})
|
||||
|
||||
refresh = create_refresh_watch_tool(
|
||||
workspace_id=db_workspace.id, thread_id=thread.id, auth_context=auth
|
||||
)
|
||||
out = await refresh.ainvoke({})
|
||||
|
||||
assert out["status"] == "refreshing"
|
||||
assert out["refreshed_ids"] == [started["automation_id"]]
|
||||
assert len(enqueue_spy) == 1
|
||||
|
||||
|
||||
async def test_tools_error_without_thread_or_auth() -> None:
|
||||
tool = create_start_watch_tool(
|
||||
workspace_id=1, thread_id=None, auth_context=None
|
||||
)
|
||||
out = await tool.ainvoke({"message": "m", "cron": _CRON, "timezone": _TZ})
|
||||
assert out["status"] == "error"
|
||||
|
|
@ -0,0 +1,118 @@
|
|||
"""Unit tests for the ``start_watch`` tool on the intelligence_agent.
|
||||
|
||||
``start_watch`` binds a recurring watch to the *current* chat: it distils the
|
||||
question + cadence the agent extracted and creates a ``schedule`` +
|
||||
``chat_message`` automation via the watch service. These tests fake the watch
|
||||
service / session so no DB is touched; they pin that the tool forwards the
|
||||
current chat id + system auth and reports a clear outcome.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
class _FakeSessionCM:
|
||||
async def __aenter__(self) -> Any:
|
||||
return MagicMock(name="session")
|
||||
|
||||
async def __aexit__(self, *_exc: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _patch_deps(monkeypatch: pytest.MonkeyPatch, *, created: Any) -> AsyncMock:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools import (
|
||||
start_watch as mod,
|
||||
)
|
||||
|
||||
fake_create_watch = AsyncMock(return_value=created)
|
||||
monkeypatch.setattr(mod, "create_watch", fake_create_watch)
|
||||
monkeypatch.setattr(mod, "AutomationService", MagicMock(return_value="svc"))
|
||||
monkeypatch.setattr(mod, "async_session_maker", lambda: _FakeSessionCM())
|
||||
return fake_create_watch
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_watch_binds_watch_to_current_chat(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools import (
|
||||
start_watch as mod,
|
||||
)
|
||||
|
||||
created = MagicMock(id=123)
|
||||
created.name = "Watch: prices"
|
||||
fake_create_watch = _patch_deps(monkeypatch, created=created)
|
||||
|
||||
auth = MagicMock()
|
||||
tool = mod.create_start_watch_tool(workspace_id=3, thread_id=55, auth_context=auth)
|
||||
|
||||
out = await tool.ainvoke(
|
||||
{"message": "what changed on prices?", "cron": "0 9 * * 1-5", "timezone": "UTC"}
|
||||
)
|
||||
|
||||
assert out["status"] == "watching"
|
||||
assert out["automation_id"] == 123
|
||||
|
||||
fake_create_watch.assert_awaited_once()
|
||||
kwargs = fake_create_watch.await_args.kwargs
|
||||
assert kwargs["workspace_id"] == 3
|
||||
assert kwargs["thread_id"] == 55
|
||||
assert kwargs["message"] == "what changed on prices?"
|
||||
assert kwargs["cron"] == "0 9 * * 1-5"
|
||||
assert kwargs["timezone"] == "UTC"
|
||||
|
||||
# The service is constructed with the passed-through auth context.
|
||||
assert mod.AutomationService.call_args.kwargs["auth"] is auth
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_watch_errors_without_thread_or_auth(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools import (
|
||||
start_watch as mod,
|
||||
)
|
||||
|
||||
fake_create_watch = _patch_deps(monkeypatch, created=MagicMock(id=1))
|
||||
|
||||
no_thread = mod.create_start_watch_tool(
|
||||
workspace_id=3, thread_id=None, auth_context=MagicMock()
|
||||
)
|
||||
out_no_thread = await no_thread.ainvoke(
|
||||
{"message": "x", "cron": "0 9 * * *", "timezone": "UTC"}
|
||||
)
|
||||
assert out_no_thread["status"] == "error"
|
||||
|
||||
no_auth = mod.create_start_watch_tool(
|
||||
workspace_id=3, thread_id=55, auth_context=None
|
||||
)
|
||||
out_no_auth = await no_auth.ainvoke(
|
||||
{"message": "x", "cron": "0 9 * * *", "timezone": "UTC"}
|
||||
)
|
||||
assert out_no_auth["status"] == "error"
|
||||
|
||||
fake_create_watch.assert_not_awaited()
|
||||
|
||||
|
||||
def test_load_tools_includes_start_watch_only_when_bindable() -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools.index import (
|
||||
load_tools,
|
||||
)
|
||||
|
||||
bindable = load_tools(
|
||||
dependencies={
|
||||
"workspace_id": 3,
|
||||
"thread_id": 55,
|
||||
"auth_context": MagicMock(),
|
||||
}
|
||||
)
|
||||
assert "start_watch" in {t.name for t in bindable}
|
||||
|
||||
not_bindable = load_tools(dependencies={"workspace_id": 3})
|
||||
assert "start_watch" not in {t.name for t in not_bindable}
|
||||
|
|
@ -0,0 +1,119 @@
|
|||
"""Unit tests for the ``stop_watch`` and ``refresh_watch`` chat tools.
|
||||
|
||||
Both act on the *current* chat: they look up the watches bound to this thread
|
||||
and stop them / run them now. The watch service + session are faked so no DB is
|
||||
touched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
class _FakeSessionCM:
|
||||
async def __aenter__(self) -> Any:
|
||||
return MagicMock(name="session")
|
||||
|
||||
async def __aexit__(self, *_exc: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _watch(automation_id: int) -> Any:
|
||||
a = MagicMock()
|
||||
a.id = automation_id
|
||||
return a
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_watch_stops_every_watch_on_the_chat(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools import (
|
||||
stop_watch as mod,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(mod, "AutomationService", MagicMock(return_value="svc"))
|
||||
monkeypatch.setattr(mod, "async_session_maker", lambda: _FakeSessionCM())
|
||||
monkeypatch.setattr(
|
||||
mod, "find_watches_for_thread", AsyncMock(return_value=[_watch(1), _watch(2)])
|
||||
)
|
||||
stop_service = AsyncMock()
|
||||
monkeypatch.setattr(mod, "stop_watch_service", stop_service)
|
||||
|
||||
tool = mod.create_stop_watch_tool(
|
||||
workspace_id=3, thread_id=55, auth_context=MagicMock()
|
||||
)
|
||||
out = await tool.ainvoke({})
|
||||
|
||||
assert out["status"] == "stopped"
|
||||
assert sorted(out["stopped_ids"]) == [1, 2]
|
||||
assert stop_service.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_watch_reports_when_nothing_to_stop(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools import (
|
||||
stop_watch as mod,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(mod, "AutomationService", MagicMock(return_value="svc"))
|
||||
monkeypatch.setattr(mod, "async_session_maker", lambda: _FakeSessionCM())
|
||||
monkeypatch.setattr(mod, "find_watches_for_thread", AsyncMock(return_value=[]))
|
||||
monkeypatch.setattr(mod, "stop_watch_service", AsyncMock())
|
||||
|
||||
tool = mod.create_stop_watch_tool(
|
||||
workspace_id=3, thread_id=55, auth_context=MagicMock()
|
||||
)
|
||||
out = await tool.ainvoke({})
|
||||
|
||||
assert out["status"] == "not_watching"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_watch_runs_each_watch_now(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools import (
|
||||
refresh_watch as mod,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(mod, "AutomationService", MagicMock(return_value="svc"))
|
||||
monkeypatch.setattr(mod, "async_session_maker", lambda: _FakeSessionCM())
|
||||
monkeypatch.setattr(
|
||||
mod, "find_watches_for_thread", AsyncMock(return_value=[_watch(7)])
|
||||
)
|
||||
run_now = AsyncMock()
|
||||
monkeypatch.setattr(mod, "run_watch_now", run_now)
|
||||
|
||||
tool = mod.create_refresh_watch_tool(
|
||||
workspace_id=3, thread_id=55, auth_context=MagicMock()
|
||||
)
|
||||
out = await tool.ainvoke({})
|
||||
|
||||
assert out["status"] == "refreshing"
|
||||
assert out["refreshed_ids"] == [7]
|
||||
run_now.assert_awaited_once()
|
||||
assert run_now.await_args.kwargs["automation_id"] == 7
|
||||
|
||||
|
||||
def test_load_tools_includes_control_tools_when_bindable() -> None:
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.intelligence_agent.tools.index import (
|
||||
load_tools,
|
||||
)
|
||||
|
||||
tools = load_tools(
|
||||
dependencies={
|
||||
"workspace_id": 3,
|
||||
"thread_id": 55,
|
||||
"auth_context": MagicMock(),
|
||||
}
|
||||
)
|
||||
names = {t.name for t in tools}
|
||||
assert {"start_watch", "stop_watch", "refresh_watch"} <= names
|
||||
Loading…
Add table
Add a link
Reference in a new issue