Merge origin/main into pr-2959
Resolve the config plumbing conflicts and keep disabled skill filtering consistent for subagent prompts after syncing with main. Made-with: Cursor
This commit is contained in:
Xubin Ren
@@ -0,0 +1,914 @@
|
||||
"""Tests for auto compact (idle TTL) feature."""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.config.schema import AgentDefaults
|
||||
from nanobot.command import CommandContext
|
||||
from nanobot.providers.base import LLMResponse
|
||||
|
||||
|
||||
def _make_loop(tmp_path: Path, session_ttl_minutes: int = 15) -> AgentLoop:
|
||||
"""Create a minimal AgentLoop for testing."""
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.estimate_prompt_tokens.return_value = (10_000, "test")
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
|
||||
provider.generation.max_tokens = 4096
|
||||
loop = AgentLoop(
|
||||
bus=bus,
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
context_window_tokens=128_000,
|
||||
session_ttl_minutes=session_ttl_minutes,
|
||||
)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
return loop
|
||||
|
||||
|
||||
def _add_turns(session, turns: int, *, prefix: str = "msg") -> None:
|
||||
"""Append simple user/assistant turns to a session."""
|
||||
for i in range(turns):
|
||||
session.add_message("user", f"{prefix} user {i}")
|
||||
session.add_message("assistant", f"{prefix} assistant {i}")
|
||||
|
||||
|
||||
class TestSessionTTLConfig:
|
||||
"""Test session TTL configuration."""
|
||||
|
||||
def test_default_ttl_is_zero(self):
|
||||
"""Default TTL should be 0 (disabled)."""
|
||||
defaults = AgentDefaults()
|
||||
assert defaults.session_ttl_minutes == 0
|
||||
|
||||
def test_custom_ttl(self):
|
||||
"""Custom TTL should be stored correctly."""
|
||||
defaults = AgentDefaults(session_ttl_minutes=30)
|
||||
assert defaults.session_ttl_minutes == 30
|
||||
|
||||
def test_user_friendly_alias_is_supported(self):
|
||||
"""Config should accept idleCompactAfterMinutes as the preferred JSON key."""
|
||||
defaults = AgentDefaults.model_validate({"idleCompactAfterMinutes": 30})
|
||||
assert defaults.session_ttl_minutes == 30
|
||||
|
||||
def test_legacy_alias_is_still_supported(self):
|
||||
"""Config should still accept the old sessionTtlMinutes key for compatibility."""
|
||||
defaults = AgentDefaults.model_validate({"sessionTtlMinutes": 30})
|
||||
assert defaults.session_ttl_minutes == 30
|
||||
|
||||
def test_serializes_with_user_friendly_alias(self):
|
||||
"""Config dumps should use idleCompactAfterMinutes for JSON output."""
|
||||
defaults = AgentDefaults(session_ttl_minutes=30)
|
||||
data = defaults.model_dump(mode="json", by_alias=True)
|
||||
assert data["idleCompactAfterMinutes"] == 30
|
||||
assert "sessionTtlMinutes" not in data
|
||||
|
||||
|
||||
class TestAgentLoopTTLParam:
|
||||
"""Test that AutoCompact receives and stores session_ttl_minutes."""
|
||||
|
||||
def test_loop_stores_ttl(self, tmp_path):
|
||||
"""AutoCompact should store the TTL value."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=25)
|
||||
assert loop.auto_compact._ttl == 25
|
||||
|
||||
def test_loop_default_ttl_zero(self, tmp_path):
|
||||
"""AutoCompact default TTL should be 0 (disabled)."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=0)
|
||||
assert loop.auto_compact._ttl == 0
|
||||
|
||||
|
||||
class TestAutoCompact:
|
||||
"""Test the _archive method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_expired_boundary(self, tmp_path):
|
||||
"""Exactly at TTL boundary should be expired (>= not >)."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
ts = datetime.now() - timedelta(minutes=15)
|
||||
assert loop.auto_compact._is_expired(ts) is True
|
||||
ts2 = datetime.now() - timedelta(minutes=14, seconds=59)
|
||||
assert loop.auto_compact._is_expired(ts2) is False
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_expired_string_timestamp(self, tmp_path):
|
||||
"""_is_expired should parse ISO string timestamps."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
ts = (datetime.now() - timedelta(minutes=20)).isoformat()
|
||||
assert loop.auto_compact._is_expired(ts) is True
|
||||
assert loop.auto_compact._is_expired(None) is False
|
||||
assert loop.auto_compact._is_expired("") is False
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_expired_only_archives_expired_sessions(self, tmp_path):
|
||||
"""With multiple sessions, only the expired one should be archived."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
# Expired session
|
||||
s1 = loop.sessions.get_or_create("cli:expired")
|
||||
s1.add_message("user", "old")
|
||||
s1.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(s1)
|
||||
# Active session
|
||||
s2 = loop.sessions.get_or_create("cli:active")
|
||||
s2.add_message("user", "recent")
|
||||
loop.sessions.save(s2)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
loop.auto_compact.check_expired(loop._schedule_background)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
active_after = loop.sessions.get_or_create("cli:active")
|
||||
assert len(active_after.messages) == 1
|
||||
assert active_after.messages[0]["content"] == "recent"
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_archives_prefix_and_keeps_recent_suffix(self, tmp_path):
|
||||
"""_archive should summarize the old prefix and keep a recent legal suffix."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archived_messages = []
|
||||
|
||||
async def _fake_archive(messages):
|
||||
archived_messages.extend(messages)
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
assert len(archived_messages) == 4
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
|
||||
assert session_after.messages[0]["content"] == "msg user 2"
|
||||
assert session_after.messages[-1]["content"] == "msg assistant 5"
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_stores_summary(self, tmp_path):
|
||||
"""_archive should store the summary in _summaries."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="hello")
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "User said hello."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
entry = loop.auto_compact._summaries.get("cli:test")
|
||||
assert entry is not None
|
||||
assert entry[0] == "User said hello."
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_empty_session(self, tmp_path):
|
||||
"""_archive on empty session should not archive."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
|
||||
archive_called = False
|
||||
|
||||
async def _fake_archive(messages):
|
||||
nonlocal archive_called
|
||||
archive_called = True
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
assert not archive_called
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == 0
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_respects_last_consolidated(self, tmp_path):
|
||||
"""_archive should only archive un-consolidated messages."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 14)
|
||||
session.last_consolidated = 18
|
||||
loop.sessions.save(session)
|
||||
|
||||
archived_count = 0
|
||||
|
||||
async def _fake_archive(messages):
|
||||
nonlocal archived_count
|
||||
archived_count = len(messages)
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
assert archived_count == 2
|
||||
await loop.close_mcp()
|
||||
|
||||
|
||||
class TestAutoCompactIdleDetection:
|
||||
"""Test idle detection triggers auto-new in _process_message."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_auto_compact_when_ttl_disabled(self, tmp_path):
|
||||
"""No auto-new should happen when TTL is 0 (disabled)."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=0)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.add_message("user", "old message")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=30)
|
||||
loop.sessions.save(session)
|
||||
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert any(m["content"] == "old message" for m in session_after.messages)
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_triggers_on_idle(self, tmp_path):
|
||||
"""Proactive auto-new archives expired session; _process_message reloads it."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="old")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archived_messages = []
|
||||
|
||||
async def _fake_archive(messages):
|
||||
archived_messages.extend(messages)
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# Simulate proactive archive completing before message arrives
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(archived_messages) == 4
|
||||
assert not any(m["content"] == "old user 0" for m in session_after.messages)
|
||||
assert any(m["content"] == "new msg" for m in session_after.messages)
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_auto_compact_when_active(self, tmp_path):
|
||||
"""No auto-new should happen when session is recently active."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.add_message("user", "recent message")
|
||||
loop.sessions.save(session)
|
||||
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert any(m["content"] == "recent message" for m in session_after.messages)
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_does_not_affect_priority_commands(self, tmp_path):
|
||||
"""Priority commands (/stop, /restart) bypass _process_message entirely via run()."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.add_message("user", "old message")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
# Priority commands are dispatched in run() before _process_message is called.
|
||||
# Simulate that path directly via dispatch_priority.
|
||||
raw = "/stop"
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content=raw)
|
||||
ctx = CommandContext(msg=msg, session=session, key="cli:test", raw=raw, loop=loop)
|
||||
result = await loop.commands.dispatch_priority(ctx)
|
||||
assert result is not None
|
||||
assert "stopped" in result.content.lower() or "no active task" in result.content.lower()
|
||||
|
||||
# Session should be untouched since priority commands skip _process_message
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert any(m["content"] == "old message" for m in session_after.messages)
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_with_slash_new(self, tmp_path):
|
||||
"""Auto-new fires before /new dispatches; session is cleared twice but idempotent."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
for i in range(4):
|
||||
session.add_message("user", f"msg{i}")
|
||||
session.add_message("assistant", f"resp{i}")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
||||
response = await loop._process_message(msg)
|
||||
|
||||
assert response is not None
|
||||
assert "new session started" in response.content.lower()
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
# Session is empty (auto-new archived and cleared, /new cleared again)
|
||||
assert len(session_after.messages) == 0
|
||||
await loop.close_mcp()
|
||||
|
||||
|
||||
class TestAutoCompactSystemMessages:
|
||||
"""Test that auto-new also works for system messages."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_triggers_for_system_messages(self, tmp_path):
|
||||
"""Proactive auto-new archives expired session; system messages reload it."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="old")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# Simulate proactive archive completing before system message arrives
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="subagent", chat_id="cli:test",
|
||||
content="subagent result",
|
||||
)
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert not any(
|
||||
m["content"] == "old user 0"
|
||||
for m in session_after.messages
|
||||
)
|
||||
await loop.close_mcp()
|
||||
|
||||
|
||||
class TestAutoCompactEdgeCases:
|
||||
"""Edge cases for auto session new."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_with_nothing_summary(self, tmp_path):
|
||||
"""Auto-new should not inject when archive produces '(nothing)'."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="thanks")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
loop.provider.chat_with_retry = AsyncMock(
|
||||
return_value=LLMResponse(content="(nothing)", tool_calls=[])
|
||||
)
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
|
||||
# "(nothing)" summary should not be stored
|
||||
assert "cli:test" not in loop.auto_compact._summaries
|
||||
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_archive_failure_still_keeps_recent_suffix(self, tmp_path):
|
||||
"""Auto-new should keep the recent suffix even if LLM archive falls back to raw dump."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="important")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
loop.provider.chat_with_retry = AsyncMock(side_effect=Exception("API down"))
|
||||
|
||||
# Should not raise
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
|
||||
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_compact_preserves_runtime_checkpoint_before_check(self, tmp_path):
|
||||
"""Short expired sessions keep recent messages; checkpoint restore still works on resume."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.metadata[AgentLoop._RUNTIME_CHECKPOINT_KEY] = {
|
||||
"assistant_message": {"role": "assistant", "content": "interrupted response"},
|
||||
"completed_tool_results": [],
|
||||
"pending_tool_calls": [],
|
||||
}
|
||||
session.add_message("user", "previous message")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archived_messages = []
|
||||
|
||||
async def _fake_archive(messages):
|
||||
archived_messages.extend(messages)
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# Simulate proactive archive completing before message arrives
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="continue")
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert archived_messages == []
|
||||
assert any(m["content"] == "previous message" for m in session_after.messages)
|
||||
assert any(m["content"] == "interrupted response" for m in session_after.messages)
|
||||
|
||||
await loop.close_mcp()
|
||||
|
||||
|
||||
class TestAutoCompactIntegration:
|
||||
"""End-to-end test of auto session new feature."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_lifecycle(self, tmp_path):
|
||||
"""
|
||||
Full lifecycle: messages -> idle -> auto-new -> archive -> clear -> summary injected as runtime context.
|
||||
"""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
|
||||
# Phase 1: User has a conversation longer than the retained recent suffix
|
||||
session.add_message("user", "I'm learning English, teach me past tense")
|
||||
session.add_message("assistant", "Past tense is used for actions completed in the past...")
|
||||
session.add_message("user", "Give me an example")
|
||||
session.add_message("assistant", '"I walked to the store yesterday."')
|
||||
session.add_message("user", "Give me another example")
|
||||
session.add_message("assistant", '"She visited Paris last year."')
|
||||
session.add_message("user", "Quiz me")
|
||||
session.add_message("assistant", "What is the past tense of go?")
|
||||
session.add_message("user", "I think it is went")
|
||||
session.add_message("assistant", "Correct.")
|
||||
loop.sessions.save(session)
|
||||
|
||||
# Phase 2: Time passes (simulate idle)
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
# Phase 3: User returns with a new message
|
||||
loop.provider.chat_with_retry = AsyncMock(
|
||||
return_value=LLMResponse(
|
||||
content="User is learning English past tense. Example: 'I walked to the store yesterday.'",
|
||||
tool_calls=[],
|
||||
)
|
||||
)
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="cli", sender_id="user", chat_id="test",
|
||||
content="Let's continue, teach me present perfect",
|
||||
)
|
||||
response = await loop._process_message(msg)
|
||||
|
||||
# Phase 4: Verify
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
|
||||
# The oldest messages should be trimmed from live session history
|
||||
assert not any(
|
||||
"past tense is used" in str(m.get("content", "")) for m in session_after.messages
|
||||
)
|
||||
|
||||
# Summary should NOT be persisted in session (ephemeral, one-shot)
|
||||
assert not any(
|
||||
"[Resumed Session]" in str(m.get("content", "")) for m in session_after.messages
|
||||
)
|
||||
# Runtime context end marker should NOT be persisted
|
||||
assert not any(
|
||||
"[/Runtime Context]" in str(m.get("content", "")) for m in session_after.messages
|
||||
)
|
||||
|
||||
# Pending summary should be consumed (one-shot)
|
||||
assert "cli:test" not in loop.auto_compact._summaries
|
||||
|
||||
# The new message should be processed (response exists)
|
||||
assert response is not None
|
||||
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_context_markers_not_persisted_for_multi_paragraph_turn(self, tmp_path):
|
||||
"""Auto-compact resume context must not leak runtime markers into persisted session history."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.add_message("user", "old message")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# Simulate proactive archive completing before message arrives
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="cli", sender_id="user", chat_id="test",
|
||||
content="Paragraph one\n\nParagraph two\n\nParagraph three",
|
||||
)
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert any(m.get("content") == "old message" for m in session_after.messages)
|
||||
for persisted in session_after.messages:
|
||||
content = str(persisted.get("content", ""))
|
||||
assert "[Runtime Context" not in content
|
||||
assert "[/Runtime Context]" not in content
|
||||
await loop.close_mcp()
|
||||
|
||||
|
||||
class TestProactiveAutoCompact:
|
||||
"""Test proactive auto-new on idle ticks (TimeoutError path in run loop)."""
|
||||
|
||||
@staticmethod
|
||||
async def _run_check_expired(loop):
|
||||
"""Helper: run check_expired via callback and wait for background tasks."""
|
||||
loop.auto_compact.check_expired(loop._schedule_background)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_check_when_ttl_disabled(self, tmp_path):
|
||||
"""check_expired should be a no-op when TTL is 0."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=0)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.add_message("user", "old message")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=30)
|
||||
loop.sessions.save(session)
|
||||
|
||||
await self._run_check_expired(loop)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == 1
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proactive_archive_on_idle_tick(self, tmp_path):
|
||||
"""Expired session should be archived during idle tick."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 5, prefix="old")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archived_messages = []
|
||||
|
||||
async def _fake_archive(messages):
|
||||
archived_messages.extend(messages)
|
||||
return "User chatted about old things."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await self._run_check_expired(loop)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
|
||||
assert len(archived_messages) == 2
|
||||
entry = loop.auto_compact._summaries.get("cli:test")
|
||||
assert entry is not None
|
||||
assert entry[0] == "User chatted about old things."
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_proactive_archive_when_active(self, tmp_path):
|
||||
"""Recently active session should NOT be archived on idle tick."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.add_message("user", "recent message")
|
||||
loop.sessions.save(session)
|
||||
|
||||
await self._run_check_expired(loop)
|
||||
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == 1
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_duplicate_archive(self, tmp_path):
|
||||
"""Should not archive the same session twice if already in progress."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="old")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archive_count = 0
|
||||
started = asyncio.Event()
|
||||
block_forever = asyncio.Event()
|
||||
|
||||
async def _slow_archive(messages):
|
||||
nonlocal archive_count
|
||||
archive_count += 1
|
||||
started.set()
|
||||
await block_forever.wait()
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _slow_archive
|
||||
|
||||
# First call starts archiving via callback
|
||||
loop.auto_compact.check_expired(loop._schedule_background)
|
||||
await started.wait()
|
||||
assert archive_count == 1
|
||||
|
||||
# Second call should skip (key is in _archiving)
|
||||
loop.auto_compact.check_expired(loop._schedule_background)
|
||||
await asyncio.sleep(0.05)
|
||||
assert archive_count == 1
|
||||
|
||||
# Clean up
|
||||
block_forever.set()
|
||||
await asyncio.sleep(0.1)
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proactive_archive_error_does_not_block(self, tmp_path):
|
||||
"""Proactive archive failure should be caught and not block future ticks."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="old")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _failing_archive(messages):
|
||||
raise RuntimeError("LLM down")
|
||||
|
||||
loop.consolidator.archive = _failing_archive
|
||||
|
||||
# Should not raise
|
||||
await self._run_check_expired(loop)
|
||||
|
||||
# Key should be removed from _archiving (finally block)
|
||||
assert "cli:test" not in loop.auto_compact._archiving
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proactive_archive_skips_empty_sessions(self, tmp_path):
|
||||
"""Proactive archive should not call LLM for sessions with no un-consolidated messages."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archive_called = False
|
||||
|
||||
async def _fake_archive(messages):
|
||||
nonlocal archive_called
|
||||
archive_called = True
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await self._run_check_expired(loop)
|
||||
|
||||
assert not archive_called
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_reschedule_after_successful_archive(self, tmp_path):
|
||||
"""Already-archived session should NOT be re-scheduled on subsequent ticks."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 5, prefix="old")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archive_count = 0
|
||||
|
||||
async def _fake_archive(messages):
|
||||
nonlocal archive_count
|
||||
archive_count += 1
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# First tick: archives the session
|
||||
await self._run_check_expired(loop)
|
||||
assert archive_count == 1
|
||||
|
||||
# Second tick: should NOT re-schedule (updated_at is fresh after clear)
|
||||
await self._run_check_expired(loop)
|
||||
assert archive_count == 1 # Still 1, not re-scheduled
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_skip_refreshes_updated_at_prevents_reschedule(self, tmp_path):
|
||||
"""Empty session skip refreshes updated_at, preventing immediate re-scheduling."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archive_count = 0
|
||||
|
||||
async def _fake_archive(messages):
|
||||
nonlocal archive_count
|
||||
archive_count += 1
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# First tick: skips (no messages), refreshes updated_at
|
||||
await self._run_check_expired(loop)
|
||||
assert archive_count == 0
|
||||
|
||||
# Second tick: should NOT re-schedule because updated_at is fresh
|
||||
await self._run_check_expired(loop)
|
||||
assert archive_count == 0
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_can_be_compacted_again_after_new_messages(self, tmp_path):
|
||||
"""After successful compact + user sends new messages + idle again, should compact again."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 5, prefix="first")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
archive_count = 0
|
||||
|
||||
async def _fake_archive(messages):
|
||||
nonlocal archive_count
|
||||
archive_count += 1
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# First compact cycle
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
assert archive_count == 1
|
||||
|
||||
# User returns, sends new messages
|
||||
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="second topic")
|
||||
await loop._process_message(msg)
|
||||
|
||||
# Simulate idle again
|
||||
loop.sessions.invalidate("cli:test")
|
||||
session2 = loop.sessions.get_or_create("cli:test")
|
||||
session2.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session2)
|
||||
|
||||
# Second compact cycle should succeed
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
assert archive_count == 2
|
||||
await loop.close_mcp()
|
||||
|
||||
|
||||
class TestSummaryPersistence:
|
||||
"""Test that summary survives restart via session metadata."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summary_persisted_in_session_metadata(self, tmp_path):
|
||||
"""After archive, _last_summary should be in session metadata."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="hello")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "User said hello."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
# Summary should be persisted in session metadata
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
meta = session_after.metadata.get("_last_summary")
|
||||
assert meta is not None
|
||||
assert meta["text"] == "User said hello."
|
||||
assert "last_active" in meta
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summary_recovered_after_restart(self, tmp_path):
|
||||
"""Summary should be recovered from metadata when _summaries is empty (simulates restart)."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="hello")
|
||||
last_active = datetime.now() - timedelta(minutes=20)
|
||||
session.updated_at = last_active
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "User said hello."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
# Archive
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
# Simulate restart: clear in-memory state
|
||||
loop.auto_compact._summaries.clear()
|
||||
loop.sessions.invalidate("cli:test")
|
||||
|
||||
# prepare_session should recover summary from metadata
|
||||
reloaded = loop.sessions.get_or_create("cli:test")
|
||||
assert len(reloaded.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
|
||||
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
|
||||
assert summary is not None
|
||||
assert "User said hello." in summary
|
||||
assert "Inactive for" in summary
|
||||
# Metadata should be cleaned up after consumption
|
||||
assert "_last_summary" not in reloaded.metadata
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_cleanup_no_leak(self, tmp_path):
|
||||
"""_last_summary should be removed from metadata after being consumed."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="hello")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
# Clear in-memory to force metadata path
|
||||
loop.auto_compact._summaries.clear()
|
||||
loop.sessions.invalidate("cli:test")
|
||||
reloaded = loop.sessions.get_or_create("cli:test")
|
||||
|
||||
# First call: consumes from metadata
|
||||
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
assert summary is not None
|
||||
|
||||
# Second call: no summary (already consumed)
|
||||
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
assert summary2 is None
|
||||
assert "_last_summary" not in reloaded.metadata
|
||||
await loop.close_mcp()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_cleanup_on_inmemory_path(self, tmp_path):
|
||||
"""In-memory _summaries path should also clean up _last_summary from metadata."""
|
||||
loop = _make_loop(tmp_path, session_ttl_minutes=15)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
_add_turns(session, 6, prefix="hello")
|
||||
session.updated_at = datetime.now() - timedelta(minutes=20)
|
||||
loop.sessions.save(session)
|
||||
|
||||
async def _fake_archive(messages):
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_archive
|
||||
|
||||
await loop.auto_compact._archive("cli:test")
|
||||
|
||||
# Both _summaries and metadata have the summary
|
||||
assert "cli:test" in loop.auto_compact._summaries
|
||||
loop.sessions.invalidate("cli:test")
|
||||
reloaded = loop.sessions.get_or_create("cli:test")
|
||||
assert "_last_summary" in reloaded.metadata
|
||||
|
||||
# In-memory path is taken (no restart)
|
||||
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
assert summary is not None
|
||||
# Metadata should also be cleaned up
|
||||
assert "_last_summary" not in reloaded.metadata
|
||||
await loop.close_mcp()
|
||||
@@ -46,7 +46,7 @@ class TestConsolidatorSummarize:
|
||||
{"role": "assistant", "content": "Done, fixed the race condition."},
|
||||
]
|
||||
result = await consolidator.archive(messages)
|
||||
assert result is True
|
||||
assert result == "User fixed a bug in the auth module."
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
|
||||
@@ -55,14 +55,14 @@ class TestConsolidatorSummarize:
|
||||
mock_provider.chat_with_retry.side_effect = Exception("API error")
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
result = await consolidator.archive(messages)
|
||||
assert result is True # always succeeds
|
||||
assert result is None # no summary on raw dump fallback
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
assert "[RAW]" in entries[0]["content"]
|
||||
|
||||
async def test_summarize_skips_empty_messages(self, consolidator):
|
||||
result = await consolidator.archive([])
|
||||
assert result is False
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestConsolidatorTokenBudget:
|
||||
@@ -76,3 +76,52 @@ class TestConsolidatorTokenBudget:
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
consolidator.archive.assert_not_called()
|
||||
|
||||
async def test_chunk_cap_preserves_user_turn_boundary(self, consolidator):
|
||||
"""Chunk cap should rewind to the last user boundary within the cap."""
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
session = MagicMock()
|
||||
session.last_consolidated = 0
|
||||
session.key = "test:key"
|
||||
session.messages = [
|
||||
{
|
||||
"role": "user" if i in {0, 50, 61} else "assistant",
|
||||
"content": f"m{i}",
|
||||
}
|
||||
for i in range(70)
|
||||
]
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(
|
||||
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
|
||||
)
|
||||
consolidator.pick_consolidation_boundary = MagicMock(return_value=(61, 999))
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
|
||||
archived_chunk = consolidator.archive.await_args.args[0]
|
||||
assert len(archived_chunk) == 50
|
||||
assert archived_chunk[0]["content"] == "m0"
|
||||
assert archived_chunk[-1]["content"] == "m49"
|
||||
assert session.last_consolidated == 50
|
||||
|
||||
async def test_chunk_cap_skips_when_no_user_boundary_within_cap(self, consolidator):
|
||||
"""If the cap would cut mid-turn, consolidation should skip that round."""
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
session = MagicMock()
|
||||
session.last_consolidated = 0
|
||||
session.key = "test:key"
|
||||
session.messages = [
|
||||
{
|
||||
"role": "user" if i in {0, 61} else "assistant",
|
||||
"content": f"m{i}",
|
||||
}
|
||||
for i in range(70)
|
||||
]
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(1200, "tiktoken"))
|
||||
consolidator.pick_consolidation_boundary = MagicMock(return_value=(61, 999))
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
|
||||
consolidator.archive.assert_not_awaited()
|
||||
assert session.last_consolidated == 0
|
||||
|
||||
@@ -307,7 +307,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
|
||||
)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
|
||||
content, tools_used, messages = await loop._run_agent_loop(
|
||||
content, tools_used, messages, _, _ = await loop._run_agent_loop(
|
||||
[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
@@ -331,7 +331,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
|
||||
)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
|
||||
content, _, _ = await loop._run_agent_loop(
|
||||
content, _, _, _, _ = await loop._run_agent_loop(
|
||||
[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
@@ -373,7 +373,7 @@ async def test_agent_loop_no_hooks_backward_compat(tmp_path):
|
||||
loop.tools.execute = AsyncMock(return_value="ok")
|
||||
loop.max_iterations = 2
|
||||
|
||||
content, tools_used, _ = await loop._run_agent_loop([])
|
||||
content, tools_used, _, _, _ = await loop._run_agent_loop([])
|
||||
assert content == (
|
||||
"I reached the maximum number of tool call iterations (2) "
|
||||
"without completing the task. You can try breaking the task into smaller steps."
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Tests for MCP connection lifecycle in AgentLoop."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.queue import MessageBus
|
||||
|
||||
|
||||
def _make_loop(tmp_path, *, mcp_servers: dict | None = None) -> AgentLoop:
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.generation.max_tokens = 4096
|
||||
return AgentLoop(
|
||||
bus=bus,
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
mcp_servers=mcp_servers or {"test": object()},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch: pytest.MonkeyPatch):
|
||||
loop = _make_loop(tmp_path)
|
||||
attempts = 0
|
||||
|
||||
async def _fake_connect(_servers, _registry):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
|
||||
|
||||
await loop._connect_mcp()
|
||||
await loop._connect_mcp()
|
||||
|
||||
assert attempts == 2
|
||||
assert loop._mcp_connected is False
|
||||
assert loop._mcp_stacks == {}
|
||||
+1161
-3
File diff suppressed because it is too large
Load Diff
@@ -5,11 +5,17 @@ from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
discord = pytest.importorskip("discord")
|
||||
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.discord import MAX_MESSAGE_LEN, DiscordBotClient, DiscordChannel, DiscordConfig
|
||||
from nanobot.channels.discord import (
|
||||
MAX_MESSAGE_LEN,
|
||||
DiscordBotClient,
|
||||
DiscordChannel,
|
||||
DiscordConfig,
|
||||
)
|
||||
from nanobot.command.builtin import build_help_text
|
||||
|
||||
|
||||
@@ -18,9 +24,11 @@ class _FakeDiscordClient:
|
||||
instances: list["_FakeDiscordClient"] = []
|
||||
start_error: Exception | None = None
|
||||
|
||||
def __init__(self, owner, *, intents) -> None:
|
||||
def __init__(self, owner, *, intents, proxy=None, proxy_auth=None) -> None:
|
||||
self.owner = owner
|
||||
self.intents = intents
|
||||
self.proxy = proxy
|
||||
self.proxy_auth = proxy_auth
|
||||
self.closed = False
|
||||
self.ready = True
|
||||
self.channels: dict[int, object] = {}
|
||||
@@ -53,7 +61,9 @@ class _FakeDiscordClient:
|
||||
|
||||
class _FakeAttachment:
|
||||
# Attachment double that can simulate successful or failing save() calls.
|
||||
def __init__(self, attachment_id: int, filename: str, *, size: int = 1, fail: bool = False) -> None:
|
||||
def __init__(
|
||||
self, attachment_id: int, filename: str, *, size: int = 1, fail: bool = False
|
||||
) -> None:
|
||||
self.id = attachment_id
|
||||
self.filename = filename
|
||||
self.size = size
|
||||
@@ -211,7 +221,7 @@ async def test_start_handles_client_construction_failure(monkeypatch) -> None:
|
||||
MessageBus(),
|
||||
)
|
||||
|
||||
def _boom(owner, *, intents):
|
||||
def _boom(owner, *, intents, proxy=None, proxy_auth=None):
|
||||
raise RuntimeError("bad client")
|
||||
|
||||
monkeypatch.setattr("nanobot.channels.discord.DiscordBotClient", _boom)
|
||||
@@ -514,9 +524,7 @@ async def test_slash_new_forwards_when_user_is_allowlisted() -> None:
|
||||
assert new_cmd is not None
|
||||
await new_cmd.callback(interaction)
|
||||
|
||||
assert interaction.response.messages == [
|
||||
{"content": "Processing /new...", "ephemeral": True}
|
||||
]
|
||||
assert interaction.response.messages == [{"content": "Processing /new...", "ephemeral": True}]
|
||||
assert len(handled) == 1
|
||||
assert handled[0]["content"] == "/new"
|
||||
assert handled[0]["sender_id"] == "123"
|
||||
@@ -590,9 +598,7 @@ async def test_slash_help_returns_ephemeral_help_text() -> None:
|
||||
assert help_cmd is not None
|
||||
await help_cmd.callback(interaction)
|
||||
|
||||
assert interaction.response.messages == [
|
||||
{"content": build_help_text(), "ephemeral": True}
|
||||
]
|
||||
assert interaction.response.messages == [{"content": build_help_text(), "ephemeral": True}]
|
||||
assert handled == []
|
||||
|
||||
|
||||
@@ -727,11 +733,13 @@ async def test_start_typing_uses_typing_context_when_trigger_typing_missing() ->
|
||||
def typing(self):
|
||||
async def _waiter():
|
||||
await release.wait()
|
||||
|
||||
# Hold the loop so task remains active until explicitly stopped.
|
||||
class _Ctx(_TypingCtx):
|
||||
async def __aenter__(self):
|
||||
await super().__aenter__()
|
||||
await _waiter()
|
||||
|
||||
return _Ctx()
|
||||
|
||||
typing_channel = _NoTriggerChannel(channel_id=123)
|
||||
@@ -745,3 +753,117 @@ async def test_start_typing_uses_typing_context_when_trigger_typing_missing() ->
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert channel._typing_tasks == {}
|
||||
|
||||
|
||||
def test_config_accepts_proxy_fields() -> None:
|
||||
config = DiscordConfig(
|
||||
enabled=True,
|
||||
token="token",
|
||||
allow_from=["*"],
|
||||
proxy="http://127.0.0.1:7890",
|
||||
proxy_username="user",
|
||||
proxy_password="pass",
|
||||
)
|
||||
assert config.proxy == "http://127.0.0.1:7890"
|
||||
assert config.proxy_username == "user"
|
||||
assert config.proxy_password == "pass"
|
||||
|
||||
|
||||
def test_config_proxy_defaults_to_none() -> None:
|
||||
config = DiscordConfig(enabled=True, token="token", allow_from=["*"])
|
||||
assert config.proxy is None
|
||||
assert config.proxy_username is None
|
||||
assert config.proxy_password is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_passes_proxy_to_client(monkeypatch) -> None:
|
||||
_FakeDiscordClient.instances.clear()
|
||||
channel = DiscordChannel(
|
||||
DiscordConfig(
|
||||
enabled=True,
|
||||
token="token",
|
||||
allow_from=["*"],
|
||||
proxy="http://127.0.0.1:7890",
|
||||
),
|
||||
MessageBus(),
|
||||
)
|
||||
monkeypatch.setattr("nanobot.channels.discord.DiscordBotClient", _FakeDiscordClient)
|
||||
|
||||
await channel.start()
|
||||
|
||||
assert channel.is_running is False
|
||||
assert len(_FakeDiscordClient.instances) == 1
|
||||
assert _FakeDiscordClient.instances[0].proxy == "http://127.0.0.1:7890"
|
||||
assert _FakeDiscordClient.instances[0].proxy_auth is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_passes_proxy_auth_when_credentials_provided(monkeypatch) -> None:
|
||||
aiohttp = pytest.importorskip("aiohttp")
|
||||
_FakeDiscordClient.instances.clear()
|
||||
channel = DiscordChannel(
|
||||
DiscordConfig(
|
||||
enabled=True,
|
||||
token="token",
|
||||
allow_from=["*"],
|
||||
proxy="http://127.0.0.1:7890",
|
||||
proxy_username="user",
|
||||
proxy_password="pass",
|
||||
),
|
||||
MessageBus(),
|
||||
)
|
||||
monkeypatch.setattr("nanobot.channels.discord.DiscordBotClient", _FakeDiscordClient)
|
||||
|
||||
await channel.start()
|
||||
|
||||
assert channel.is_running is False
|
||||
assert len(_FakeDiscordClient.instances) == 1
|
||||
assert _FakeDiscordClient.instances[0].proxy == "http://127.0.0.1:7890"
|
||||
assert _FakeDiscordClient.instances[0].proxy_auth is not None
|
||||
assert isinstance(_FakeDiscordClient.instances[0].proxy_auth, aiohttp.BasicAuth)
|
||||
assert _FakeDiscordClient.instances[0].proxy_auth.login == "user"
|
||||
assert _FakeDiscordClient.instances[0].proxy_auth.password == "pass"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_no_proxy_auth_when_only_username(monkeypatch) -> None:
|
||||
_FakeDiscordClient.instances.clear()
|
||||
channel = DiscordChannel(
|
||||
DiscordConfig(
|
||||
enabled=True,
|
||||
token="token",
|
||||
allow_from=["*"],
|
||||
proxy="http://127.0.0.1:7890",
|
||||
proxy_username="user",
|
||||
),
|
||||
MessageBus(),
|
||||
)
|
||||
monkeypatch.setattr("nanobot.channels.discord.DiscordBotClient", _FakeDiscordClient)
|
||||
|
||||
await channel.start()
|
||||
|
||||
assert channel.is_running is False
|
||||
assert _FakeDiscordClient.instances[0].proxy_auth is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_no_proxy_auth_when_only_password(monkeypatch) -> None:
|
||||
_FakeDiscordClient.instances.clear()
|
||||
channel = DiscordChannel(
|
||||
DiscordConfig(
|
||||
enabled=True,
|
||||
token="token",
|
||||
allow_from=["*"],
|
||||
proxy="http://127.0.0.1:7890",
|
||||
proxy_password="pass",
|
||||
),
|
||||
MessageBus(),
|
||||
)
|
||||
monkeypatch.setattr("nanobot.channels.discord.DiscordBotClient", _FakeDiscordClient)
|
||||
|
||||
await channel.start()
|
||||
|
||||
assert channel.is_running is False
|
||||
assert _FakeDiscordClient.instances[0].proxy == "http://127.0.0.1:7890"
|
||||
assert _FakeDiscordClient.instances[0].proxy_auth is None
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Tests for Feishu/Lark domain configuration."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.feishu import FeishuChannel, FeishuConfig
|
||||
|
||||
|
||||
def _make_channel(domain: str = "feishu") -> FeishuChannel:
|
||||
config = FeishuConfig(
|
||||
enabled=True,
|
||||
app_id="cli_test",
|
||||
app_secret="secret",
|
||||
allow_from=["*"],
|
||||
domain=domain,
|
||||
)
|
||||
ch = FeishuChannel(config, MessageBus())
|
||||
ch._client = MagicMock()
|
||||
ch._loop = None
|
||||
return ch
|
||||
|
||||
|
||||
class TestFeishuConfigDomain:
|
||||
def test_domain_default_is_feishu(self):
|
||||
config = FeishuConfig()
|
||||
assert config.domain == "feishu"
|
||||
|
||||
def test_domain_accepts_lark(self):
|
||||
config = FeishuConfig(domain="lark")
|
||||
assert config.domain == "lark"
|
||||
|
||||
def test_domain_accepts_feishu(self):
|
||||
config = FeishuConfig(domain="feishu")
|
||||
assert config.domain == "feishu"
|
||||
|
||||
def test_default_config_includes_domain(self):
|
||||
default_cfg = FeishuChannel.default_config()
|
||||
assert "domain" in default_cfg
|
||||
assert default_cfg["domain"] == "feishu"
|
||||
|
||||
def test_channel_persists_domain_from_config(self):
|
||||
ch = _make_channel(domain="lark")
|
||||
assert ch.config.domain == "lark"
|
||||
|
||||
def test_channel_persists_feishu_domain_from_config(self):
|
||||
ch = _make_channel(domain="feishu")
|
||||
assert ch.config.domain == "feishu"
|
||||
@@ -5,6 +5,7 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.feishu import FeishuChannel, FeishuConfig, _FeishuStreamBuf
|
||||
|
||||
@@ -203,6 +204,55 @@ class TestSendDelta:
|
||||
ch._client.cardkit.v1.card_element.content.assert_not_called()
|
||||
ch._client.im.v1.message.create.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_resuming_keeps_buffer(self):
|
||||
"""_resuming=True flushes text to card but keeps the buffer for the next segment."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
|
||||
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
|
||||
|
||||
assert "oc_chat1" in ch._stream_bufs
|
||||
buf = ch._stream_bufs["oc_chat1"]
|
||||
assert buf.card_id == "card_1"
|
||||
assert buf.sequence == 3
|
||||
ch._client.cardkit.v1.card_element.content.assert_called_once()
|
||||
ch._client.cardkit.v1.card.settings.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_resuming_then_final_end(self):
|
||||
"""Full multi-segment flow: resuming mid-turn, then final end closes the card."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Seg1", card_id="card_1", sequence=1, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
ch._client.cardkit.v1.card.settings.return_value = _mock_content_response()
|
||||
|
||||
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
|
||||
assert "oc_chat1" in ch._stream_bufs
|
||||
|
||||
ch._stream_bufs["oc_chat1"].text += " Seg2"
|
||||
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True})
|
||||
|
||||
assert "oc_chat1" not in ch._stream_bufs
|
||||
ch._client.cardkit.v1.card.settings.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_resuming_no_card_is_noop(self):
|
||||
"""_resuming with no card_id (card creation failed) is a safe no-op."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="text", card_id=None, sequence=0, last_edit=0.0,
|
||||
)
|
||||
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
|
||||
|
||||
assert "oc_chat1" in ch._stream_bufs
|
||||
ch._client.cardkit.v1.card_element.content.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_without_buf_is_noop(self):
|
||||
ch = _make_channel()
|
||||
@@ -239,6 +289,146 @@ class TestSendDelta:
|
||||
assert buf.sequence == 7
|
||||
|
||||
|
||||
class TestToolHintInlineStreaming:
|
||||
"""Tool hint messages should be inlined into active streaming cards."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_hint_inlined_when_stream_active(self):
|
||||
"""With an active streaming buffer, tool hint appends to the card."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
|
||||
msg = OutboundMessage(
|
||||
channel="feishu", chat_id="oc_chat1",
|
||||
content='web_fetch("https://example.com")',
|
||||
metadata={"_tool_hint": True},
|
||||
)
|
||||
await ch.send(msg)
|
||||
|
||||
buf = ch._stream_bufs["oc_chat1"]
|
||||
assert '🔧 web_fetch("https://example.com")' in buf.text
|
||||
assert buf.sequence == 3
|
||||
ch._client.cardkit.v1.card_element.content.assert_called_once()
|
||||
ch._client.im.v1.message.create.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_hint_preserved_on_next_delta(self):
|
||||
"""When new delta arrives, the tool hint is kept as permanent content and delta appends after it."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Partial answer\n\n🔧 web_fetch(\"url\")\n\n",
|
||||
card_id="card_1", sequence=3, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
|
||||
await ch.send_delta("oc_chat1", " continued")
|
||||
|
||||
buf = ch._stream_bufs["oc_chat1"]
|
||||
assert "Partial answer" in buf.text
|
||||
assert "🔧 web_fetch" in buf.text
|
||||
assert buf.text.endswith(" continued")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_hint_fallback_when_no_stream(self):
|
||||
"""Without an active buffer, tool hint falls back to a standalone card."""
|
||||
ch = _make_channel()
|
||||
ch._client.im.v1.message.create.return_value = _mock_send_response("om_hint")
|
||||
|
||||
msg = OutboundMessage(
|
||||
channel="feishu", chat_id="oc_chat1",
|
||||
content='read_file("path")',
|
||||
metadata={"_tool_hint": True},
|
||||
)
|
||||
await ch.send(msg)
|
||||
|
||||
assert "oc_chat1" not in ch._stream_bufs
|
||||
ch._client.im.v1.message.create.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consecutive_tool_hints_append(self):
|
||||
"""When multiple tool hints arrive consecutively, each appends to the card."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
|
||||
msg1 = OutboundMessage(
|
||||
channel="feishu", chat_id="oc_chat1",
|
||||
content='$ cd /project', metadata={"_tool_hint": True},
|
||||
)
|
||||
await ch.send(msg1)
|
||||
|
||||
msg2 = OutboundMessage(
|
||||
channel="feishu", chat_id="oc_chat1",
|
||||
content='$ git status', metadata={"_tool_hint": True},
|
||||
)
|
||||
await ch.send(msg2)
|
||||
|
||||
buf = ch._stream_bufs["oc_chat1"]
|
||||
assert "$ cd /project" in buf.text
|
||||
assert "$ git status" in buf.text
|
||||
assert buf.text.startswith("Partial answer")
|
||||
assert "🔧 $ cd /project" in buf.text
|
||||
assert "🔧 $ git status" in buf.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_hint_preserved_on_resuming_flush(self):
|
||||
"""When _resuming flushes the buffer, tool hint is kept as permanent content."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Partial answer\n\n🔧 $ cd /project\n\n",
|
||||
card_id="card_1", sequence=2, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
|
||||
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
|
||||
|
||||
buf = ch._stream_bufs["oc_chat1"]
|
||||
assert "Partial answer" in buf.text
|
||||
assert "🔧 $ cd /project" in buf.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_hint_preserved_on_final_stream_end(self):
|
||||
"""When final _stream_end closes the card, tool hint is kept in the final text."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Final content\n\n🔧 web_fetch(\"url\")\n\n",
|
||||
card_id="card_1", sequence=3, last_edit=0.0,
|
||||
)
|
||||
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
|
||||
ch._client.cardkit.v1.card.settings.return_value = _mock_content_response()
|
||||
|
||||
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True})
|
||||
|
||||
assert "oc_chat1" not in ch._stream_bufs
|
||||
update_call = ch._client.cardkit.v1.card_element.content.call_args[0][0]
|
||||
assert "🔧" in update_call.body.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_tool_hint_is_noop(self):
|
||||
"""Empty or whitespace-only tool hint content is silently ignored."""
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
|
||||
)
|
||||
|
||||
for content in ("", " ", "\t\n"):
|
||||
msg = OutboundMessage(
|
||||
channel="feishu", chat_id="oc_chat1",
|
||||
content=content, metadata={"_tool_hint": True},
|
||||
)
|
||||
await ch.send(msg)
|
||||
|
||||
buf = ch._stream_bufs["oc_chat1"]
|
||||
assert buf.text == "Partial answer"
|
||||
assert buf.sequence == 2
|
||||
ch._client.cardkit.v1.card_element.content.assert_not_called()
|
||||
|
||||
|
||||
class TestSendMessageReturnsId:
|
||||
def test_returns_message_id_on_success(self):
|
||||
ch = _make_channel()
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Tests for QQ channel media support: helpers, send, inbound, and upload."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
from nanobot.channels import qq
|
||||
|
||||
QQ_AVAILABLE = getattr(qq, "QQ_AVAILABLE", False)
|
||||
except ImportError:
|
||||
QQ_AVAILABLE = False
|
||||
|
||||
if not QQ_AVAILABLE:
|
||||
pytest.skip("QQ dependencies not installed (qq-botpy)", allow_module_level=True)
|
||||
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.qq import (
|
||||
QQ_FILE_TYPE_FILE,
|
||||
QQ_FILE_TYPE_IMAGE,
|
||||
QQChannel,
|
||||
QQConfig,
|
||||
_guess_send_file_type,
|
||||
_is_image_name,
|
||||
_sanitize_filename,
|
||||
)
|
||||
|
||||
|
||||
class _FakeApi:
|
||||
def __init__(self) -> None:
|
||||
self.c2c_calls: list[dict] = []
|
||||
self.group_calls: list[dict] = []
|
||||
|
||||
async def post_c2c_message(self, **kwargs) -> None:
|
||||
self.c2c_calls.append(kwargs)
|
||||
|
||||
async def post_group_message(self, **kwargs) -> None:
|
||||
self.group_calls.append(kwargs)
|
||||
|
||||
|
||||
class _FakeHttp:
|
||||
"""Fake _http for _post_base64file tests."""
|
||||
|
||||
def __init__(self, return_value: dict | None = None) -> None:
|
||||
self.return_value = return_value or {}
|
||||
self.calls: list[tuple] = []
|
||||
|
||||
async def request(self, route, **kwargs):
|
||||
self.calls.append((route, kwargs))
|
||||
return self.return_value
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, http_return: dict | None = None) -> None:
|
||||
self.api = _FakeApi()
|
||||
self.api._http = _FakeHttp(http_return)
|
||||
|
||||
|
||||
# ── Helper function tests (pure, no async) ──────────────────────────
|
||||
|
||||
|
||||
def test_sanitize_filename_strips_path_traversal() -> None:
|
||||
assert _sanitize_filename("../../etc/passwd") == "passwd"
|
||||
|
||||
|
||||
def test_sanitize_filename_keeps_chinese_chars() -> None:
|
||||
assert _sanitize_filename("文件（1）.jpg") == "文件（1）.jpg"
|
||||
|
||||
|
||||
def test_sanitize_filename_strips_unsafe_chars() -> None:
|
||||
result = _sanitize_filename('file<>:"|?*.txt')
|
||||
# All unsafe chars replaced with "_", but * is replaced too
|
||||
assert result.startswith("file")
|
||||
assert result.endswith(".txt")
|
||||
assert "<" not in result
|
||||
assert ">" not in result
|
||||
assert '"' not in result
|
||||
assert "|" not in result
|
||||
assert "?" not in result
|
||||
|
||||
|
||||
def test_sanitize_filename_empty_input() -> None:
|
||||
assert _sanitize_filename("") == ""
|
||||
|
||||
|
||||
def test_is_image_name_with_known_extensions() -> None:
|
||||
for ext in (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tif", ".tiff", ".ico", ".svg"):
|
||||
assert _is_image_name(f"photo{ext}") is True
|
||||
|
||||
|
||||
def test_is_image_name_with_unknown_extension() -> None:
|
||||
for ext in (".pdf", ".txt", ".mp3", ".mp4"):
|
||||
assert _is_image_name(f"doc{ext}") is False
|
||||
|
||||
|
||||
def test_guess_send_file_type_image() -> None:
|
||||
assert _guess_send_file_type("photo.png") == QQ_FILE_TYPE_IMAGE
|
||||
assert _guess_send_file_type("pic.jpg") == QQ_FILE_TYPE_IMAGE
|
||||
|
||||
|
||||
def test_guess_send_file_type_file() -> None:
|
||||
assert _guess_send_file_type("doc.pdf") == QQ_FILE_TYPE_FILE
|
||||
|
||||
|
||||
def test_guess_send_file_type_by_mime() -> None:
|
||||
# A filename with no known extension but whose mime type is image/*
|
||||
assert _guess_send_file_type("photo.xyz_image_test") == QQ_FILE_TYPE_FILE
|
||||
|
||||
|
||||
# ── send() exception handling ───────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_exception_caught_not_raised() -> None:
|
||||
"""Exceptions inside send() must not propagate."""
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
|
||||
channel._client = _FakeClient()
|
||||
|
||||
with patch.object(channel, "_send_text_only", new_callable=AsyncMock, side_effect=RuntimeError("boom")):
|
||||
await channel.send(
|
||||
OutboundMessage(channel="qq", chat_id="user1", content="hello")
|
||||
)
|
||||
# No exception raised — test passes if we get here.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_then_text() -> None:
|
||||
"""Media is sent before text when both are present."""
|
||||
import tempfile
|
||||
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
|
||||
channel._client = _FakeClient()
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
tmp = f.name
|
||||
|
||||
try:
|
||||
with patch.object(channel, "_post_base64file", new_callable=AsyncMock, return_value={"file_info": "1"}) as mock_upload:
|
||||
await channel.send(
|
||||
OutboundMessage(
|
||||
channel="qq",
|
||||
chat_id="user1",
|
||||
content="text after image",
|
||||
media=[tmp],
|
||||
metadata={"message_id": "m1"},
|
||||
)
|
||||
)
|
||||
assert mock_upload.called
|
||||
|
||||
# Text should have been sent via c2c (default chat type)
|
||||
text_calls = [c for c in channel._client.api.c2c_calls if c.get("msg_type") == 0]
|
||||
assert len(text_calls) >= 1
|
||||
assert text_calls[-1]["content"] == "text after image"
|
||||
finally:
|
||||
import os
|
||||
os.unlink(tmp)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_failure_falls_back_to_text() -> None:
|
||||
"""When _send_media returns False, a failure notice is appended."""
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
|
||||
channel._client = _FakeClient()
|
||||
|
||||
with patch.object(channel, "_send_media", new_callable=AsyncMock, return_value=False):
|
||||
await channel.send(
|
||||
OutboundMessage(
|
||||
channel="qq",
|
||||
chat_id="user1",
|
||||
content="hello",
|
||||
media=["https://example.com/bad.png"],
|
||||
metadata={"message_id": "m1"},
|
||||
)
|
||||
)
|
||||
|
||||
# Should have the failure text among the c2c calls
|
||||
failure_calls = [c for c in channel._client.api.c2c_calls if "Attachment send failed" in c.get("content", "")]
|
||||
assert len(failure_calls) == 1
|
||||
assert "bad.png" in failure_calls[0]["content"]
|
||||
|
||||
|
||||
# ── _on_message() exception handling ────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_message_exception_caught_not_raised() -> None:
|
||||
"""Missing required attributes should not crash _on_message."""
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
|
||||
channel._client = _FakeClient()
|
||||
|
||||
# Construct a message-like object that lacks 'author' — triggers AttributeError
|
||||
bad_data = SimpleNamespace(id="x1", content="hi")
|
||||
# Should not raise
|
||||
await channel._on_message(bad_data, is_group=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_message_with_attachments() -> None:
|
||||
"""Messages with attachments produce media_paths and formatted content."""
|
||||
import tempfile
|
||||
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
|
||||
channel._client = _FakeClient()
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
saved_path = f.name
|
||||
|
||||
att = SimpleNamespace(url="", filename="screenshot.png", content_type="image/png")
|
||||
|
||||
# Patch _download_to_media_dir_chunked to return the temp file path
|
||||
async def fake_download(url, filename_hint=""):
|
||||
return saved_path
|
||||
|
||||
try:
|
||||
with patch.object(channel, "_download_to_media_dir_chunked", side_effect=fake_download):
|
||||
data = SimpleNamespace(
|
||||
id="att1",
|
||||
content="look at this",
|
||||
author=SimpleNamespace(user_openid="u1"),
|
||||
attachments=[att],
|
||||
)
|
||||
await channel._on_message(data, is_group=False)
|
||||
|
||||
msg = await channel.bus.consume_inbound()
|
||||
assert "look at this" in msg.content
|
||||
assert "screenshot.png" in msg.content
|
||||
assert "Received files:" in msg.content
|
||||
assert len(msg.media) == 1
|
||||
assert msg.media[0] == saved_path
|
||||
finally:
|
||||
import os
|
||||
os.unlink(saved_path)
|
||||
|
||||
|
||||
# ── _post_base64file() ─────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_base64file_omits_file_name_for_images() -> None:
|
||||
"""file_type=1 (image) → payload must not contain file_name."""
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret"), MessageBus())
|
||||
channel._client = _FakeClient(http_return={"file_info": "img_abc"})
|
||||
|
||||
await channel._post_base64file(
|
||||
chat_id="user1",
|
||||
is_group=False,
|
||||
file_type=QQ_FILE_TYPE_IMAGE,
|
||||
file_data="ZmFrZQ==",
|
||||
file_name="photo.png",
|
||||
)
|
||||
|
||||
http = channel._client.api._http
|
||||
assert len(http.calls) == 1
|
||||
payload = http.calls[0][1]["json"]
|
||||
assert "file_name" not in payload
|
||||
assert payload["file_type"] == QQ_FILE_TYPE_IMAGE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_base64file_includes_file_name_for_files() -> None:
|
||||
"""file_type=4 (file) → payload must contain file_name."""
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret"), MessageBus())
|
||||
channel._client = _FakeClient(http_return={"file_info": "file_abc"})
|
||||
|
||||
await channel._post_base64file(
|
||||
chat_id="user1",
|
||||
is_group=False,
|
||||
file_type=QQ_FILE_TYPE_FILE,
|
||||
file_data="ZmFrZQ==",
|
||||
file_name="report.pdf",
|
||||
)
|
||||
|
||||
http = channel._client.api._http
|
||||
assert len(http.calls) == 1
|
||||
payload = http.calls[0][1]["json"]
|
||||
assert payload["file_name"] == "report.pdf"
|
||||
assert payload["file_type"] == QQ_FILE_TYPE_FILE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_base64file_filters_response_to_file_info() -> None:
|
||||
"""Response with file_info + extra fields must be filtered to only file_info."""
|
||||
channel = QQChannel(QQConfig(app_id="app", secret="secret"), MessageBus())
|
||||
channel._client = _FakeClient(http_return={
|
||||
"file_info": "fi_123",
|
||||
"file_uuid": "uuid_xxx",
|
||||
"ttl": 3600,
|
||||
})
|
||||
|
||||
result = await channel._post_base64file(
|
||||
chat_id="user1",
|
||||
is_group=False,
|
||||
file_type=QQ_FILE_TYPE_FILE,
|
||||
file_data="ZmFrZQ==",
|
||||
file_name="doc.pdf",
|
||||
)
|
||||
|
||||
assert result == {"file_info": "fi_123"}
|
||||
assert "file_uuid" not in result
|
||||
assert "ttl" not in result
|
||||
@@ -0,0 +1,598 @@
|
||||
"""Unit and lightweight integration tests for the WebSocket channel."""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import websockets
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
from websockets.frames import Close
|
||||
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.channels.websocket import (
|
||||
WebSocketChannel,
|
||||
WebSocketConfig,
|
||||
_issue_route_secret_matches,
|
||||
_normalize_config_path,
|
||||
_normalize_http_path,
|
||||
_parse_inbound_payload,
|
||||
_parse_query,
|
||||
_parse_request_path,
|
||||
)
|
||||
|
||||
# -- Shared helpers (aligned with test_websocket_integration.py) ---------------
|
||||
|
||||
_PORT = 29876
|
||||
|
||||
|
||||
def _ch(bus: Any, **kw: Any) -> WebSocketChannel:
|
||||
cfg: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"allowFrom": ["*"],
|
||||
"host": "127.0.0.1",
|
||||
"port": _PORT,
|
||||
"path": "/ws",
|
||||
"websocketRequiresToken": False,
|
||||
}
|
||||
cfg.update(kw)
|
||||
return WebSocketChannel(cfg, bus)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def bus() -> MagicMock:
|
||||
b = MagicMock()
|
||||
b.publish_inbound = AsyncMock()
|
||||
return b
|
||||
|
||||
|
||||
async def _http_get(url: str, headers: dict[str, str] | None = None) -> httpx.Response:
|
||||
"""Run GET in a thread to avoid blocking the asyncio loop shared with websockets."""
|
||||
return await asyncio.to_thread(
|
||||
functools.partial(httpx.get, url, headers=headers or {}, timeout=5.0)
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_http_path_strips_trailing_slash_except_root() -> None:
|
||||
assert _normalize_http_path("/chat/") == "/chat"
|
||||
assert _normalize_http_path("/chat?x=1") == "/chat"
|
||||
assert _normalize_http_path("/") == "/"
|
||||
|
||||
|
||||
def test_parse_request_path_matches_normalize_and_query() -> None:
|
||||
path, query = _parse_request_path("/ws/?token=secret&client_id=u1")
|
||||
assert path == _normalize_http_path("/ws/?token=secret&client_id=u1")
|
||||
assert query == _parse_query("/ws/?token=secret&client_id=u1")
|
||||
|
||||
|
||||
def test_normalize_config_path_matches_request() -> None:
|
||||
assert _normalize_config_path("/ws/") == "/ws"
|
||||
assert _normalize_config_path("/") == "/"
|
||||
|
||||
|
||||
def test_parse_query_extracts_token_and_client_id() -> None:
|
||||
query = _parse_query("/?token=secret&client_id=u1")
|
||||
assert query.get("token") == ["secret"]
|
||||
assert query.get("client_id") == ["u1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[
|
||||
("plain", "plain"),
|
||||
('{"content": "hi"}', "hi"),
|
||||
('{"text": "there"}', "there"),
|
||||
('{"message": "x"}', "x"),
|
||||
(" ", None),
|
||||
("{}", None),
|
||||
],
|
||||
)
|
||||
def test_parse_inbound_payload(raw: str, expected: str | None) -> None:
|
||||
assert _parse_inbound_payload(raw) == expected
|
||||
|
||||
|
||||
def test_parse_inbound_invalid_json_falls_back_to_raw_string() -> None:
|
||||
assert _parse_inbound_payload("{not json") == "{not json"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[
|
||||
('{"content": ""}', None), # empty string content
|
||||
('{"content": 123}', None), # non-string content
|
||||
('{"content": " "}', None), # whitespace-only content
|
||||
('["hello"]', '["hello"]'), # JSON array: not a dict, treated as plain text
|
||||
('{"unknown_key": "val"}', None), # unrecognized key
|
||||
('{"content": null}', None), # null content
|
||||
],
|
||||
)
|
||||
def test_parse_inbound_payload_edge_cases(raw: str, expected: str | None) -> None:
|
||||
assert _parse_inbound_payload(raw) == expected
|
||||
|
||||
|
||||
def test_web_socket_config_path_must_start_with_slash() -> None:
|
||||
with pytest.raises(ValueError, match='path must start with "/"'):
|
||||
WebSocketConfig(path="bad")
|
||||
|
||||
|
||||
def test_ssl_context_requires_both_cert_and_key_files() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel(
|
||||
{"enabled": True, "allowFrom": ["*"], "sslCertfile": "/tmp/c.pem", "sslKeyfile": ""},
|
||||
bus,
|
||||
)
|
||||
with pytest.raises(ValueError, match="ssl_certfile and ssl_keyfile"):
|
||||
channel._build_ssl_context()
|
||||
|
||||
|
||||
def test_default_config_includes_safe_bind_and_streaming() -> None:
|
||||
defaults = WebSocketChannel.default_config()
|
||||
assert defaults["enabled"] is False
|
||||
assert defaults["host"] == "127.0.0.1"
|
||||
assert defaults["streaming"] is True
|
||||
assert defaults["allowFrom"] == ["*"]
|
||||
assert defaults.get("tokenIssuePath", "") == ""
|
||||
|
||||
|
||||
def test_token_issue_path_must_differ_from_websocket_path() -> None:
|
||||
with pytest.raises(ValueError, match="token_issue_path must differ"):
|
||||
WebSocketConfig(path="/ws", token_issue_path="/ws")
|
||||
|
||||
|
||||
def test_issue_route_secret_matches_bearer_and_header() -> None:
|
||||
from websockets.datastructures import Headers
|
||||
|
||||
secret = "my-secret"
|
||||
bearer_headers = Headers([("Authorization", "Bearer my-secret")])
|
||||
assert _issue_route_secret_matches(bearer_headers, secret) is True
|
||||
x_headers = Headers([("X-Nanobot-Auth", "my-secret")])
|
||||
assert _issue_route_secret_matches(x_headers, secret) is True
|
||||
wrong = Headers([("Authorization", "Bearer other")])
|
||||
assert _issue_route_secret_matches(wrong, secret) is False
|
||||
|
||||
|
||||
def test_issue_route_secret_matches_empty_secret() -> None:
|
||||
from websockets.datastructures import Headers
|
||||
|
||||
# Empty secret always returns True regardless of headers
|
||||
assert _issue_route_secret_matches(Headers([]), "") is True
|
||||
assert _issue_route_secret_matches(Headers([("Authorization", "Bearer anything")]), "") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delivers_json_message_with_media_and_reply() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
|
||||
mock_ws = AsyncMock()
|
||||
channel._connections["chat-1"] = mock_ws
|
||||
|
||||
msg = OutboundMessage(
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
content="hello",
|
||||
reply_to="m1",
|
||||
media=["/tmp/a.png"],
|
||||
)
|
||||
await channel.send(msg)
|
||||
|
||||
mock_ws.send.assert_awaited_once()
|
||||
payload = json.loads(mock_ws.send.call_args[0][0])
|
||||
assert payload["event"] == "message"
|
||||
assert payload["text"] == "hello"
|
||||
assert payload["reply_to"] == "m1"
|
||||
assert payload["media"] == ["/tmp/a.png"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_missing_connection_is_noop_without_error() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
|
||||
msg = OutboundMessage(channel="websocket", chat_id="missing", content="x")
|
||||
await channel.send(msg)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_removes_connection_on_connection_closed() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True)
|
||||
channel._connections["chat-1"] = mock_ws
|
||||
|
||||
msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello")
|
||||
await channel.send(msg)
|
||||
|
||||
assert "chat-1" not in channel._connections
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_removes_connection_on_connection_closed() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True)
|
||||
channel._connections["chat-1"] = mock_ws
|
||||
|
||||
await channel.send_delta("chat-1", "chunk", {"_stream_delta": True, "_stream_id": "s1"})
|
||||
|
||||
assert "chat-1" not in channel._connections
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_emits_delta_and_stream_end() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
|
||||
mock_ws = AsyncMock()
|
||||
channel._connections["chat-1"] = mock_ws
|
||||
|
||||
await channel.send_delta("chat-1", "part", {"_stream_delta": True, "_stream_id": "sid"})
|
||||
await channel.send_delta("chat-1", "", {"_stream_end": True, "_stream_id": "sid"})
|
||||
|
||||
assert mock_ws.send.await_count == 2
|
||||
first = json.loads(mock_ws.send.call_args_list[0][0][0])
|
||||
second = json.loads(mock_ws.send.call_args_list[1][0][0])
|
||||
assert first["event"] == "delta"
|
||||
assert first["text"] == "part"
|
||||
assert first["stream_id"] == "sid"
|
||||
assert second["event"] == "stream_end"
|
||||
assert second["stream_id"] == "sid"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_non_connection_closed_exception_is_raised() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.send.side_effect = RuntimeError("unexpected")
|
||||
channel._connections["chat-1"] = mock_ws
|
||||
|
||||
msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello")
|
||||
with pytest.raises(RuntimeError, match="unexpected"):
|
||||
await channel.send(msg)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_missing_connection_is_noop() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
|
||||
# No exception, no error — just a no-op
|
||||
await channel.send_delta("nonexistent", "chunk", {"_stream_delta": True, "_stream_id": "s1"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_is_idempotent() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
|
||||
# stop() before start() should not raise
|
||||
await channel.stop()
|
||||
await channel.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_to_end_client_receives_ready_and_agent_sees_inbound(bus: MagicMock) -> None:
|
||||
port = 29876
|
||||
channel = _ch(bus, port=port)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=tester") as client:
|
||||
ready_raw = await client.recv()
|
||||
ready = json.loads(ready_raw)
|
||||
assert ready["event"] == "ready"
|
||||
assert ready["client_id"] == "tester"
|
||||
chat_id = ready["chat_id"]
|
||||
|
||||
await client.send(json.dumps({"content": "ping from client"}))
|
||||
await asyncio.sleep(0.08)
|
||||
|
||||
bus.publish_inbound.assert_awaited()
|
||||
inbound = bus.publish_inbound.call_args[0][0]
|
||||
assert inbound.channel == "websocket"
|
||||
assert inbound.sender_id == "tester"
|
||||
assert inbound.chat_id == chat_id
|
||||
assert inbound.content == "ping from client"
|
||||
|
||||
await client.send("plain text frame")
|
||||
await asyncio.sleep(0.08)
|
||||
assert bus.publish_inbound.await_count >= 2
|
||||
second = [c[0][0] for c in bus.publish_inbound.call_args_list][-1]
|
||||
assert second.content == "plain text frame"
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_rejects_handshake_when_mismatch(bus: MagicMock) -> None:
|
||||
port = 29877
|
||||
channel = _ch(bus, port=port, path="/", token="secret")
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as excinfo:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/?token=wrong"):
|
||||
pass
|
||||
assert excinfo.value.response.status_code == 401
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wrong_path_returns_404(bus: MagicMock) -> None:
|
||||
port = 29878
|
||||
channel = _ch(bus, port=port)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as excinfo:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/other"):
|
||||
pass
|
||||
assert excinfo.value.response.status_code == 404
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
def test_registry_discovers_websocket_channel() -> None:
|
||||
from nanobot.channels.registry import load_channel_class
|
||||
|
||||
cls = load_channel_class("websocket")
|
||||
assert cls.name == "websocket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_route_issues_token_then_websocket_requires_it(bus: MagicMock) -> None:
|
||||
port = 29879
|
||||
channel = _ch(
|
||||
bus, port=port,
|
||||
tokenIssuePath="/auth/token",
|
||||
tokenIssueSecret="route-secret",
|
||||
websocketRequiresToken=True,
|
||||
)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
deny = await _http_get(f"http://127.0.0.1:{port}/auth/token")
|
||||
assert deny.status_code == 401
|
||||
|
||||
issue = await _http_get(
|
||||
f"http://127.0.0.1:{port}/auth/token",
|
||||
headers={"Authorization": "Bearer route-secret"},
|
||||
)
|
||||
assert issue.status_code == 200
|
||||
token = issue.json()["token"]
|
||||
assert token.startswith("nbwt_")
|
||||
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as missing_token:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=x"):
|
||||
pass
|
||||
assert missing_token.value.response.status_code == 401
|
||||
|
||||
uri = f"ws://127.0.0.1:{port}/ws?token={token}&client_id=caller"
|
||||
async with websockets.connect(uri) as client:
|
||||
ready = json.loads(await client.recv())
|
||||
assert ready["event"] == "ready"
|
||||
assert ready["client_id"] == "caller"
|
||||
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as reuse:
|
||||
async with websockets.connect(uri):
|
||||
pass
|
||||
assert reuse.value.response.status_code == 401
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_to_end_server_pushes_streaming_deltas_to_client(bus: MagicMock) -> None:
|
||||
port = 29880
|
||||
channel = _ch(bus, port=port, streaming=True)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=stream-tester") as client:
|
||||
ready_raw = await client.recv()
|
||||
ready = json.loads(ready_raw)
|
||||
chat_id = ready["chat_id"]
|
||||
|
||||
# Server pushes deltas directly
|
||||
await channel.send_delta(
|
||||
chat_id, "Hello ", {"_stream_delta": True, "_stream_id": "s1"}
|
||||
)
|
||||
await channel.send_delta(
|
||||
chat_id, "world", {"_stream_delta": True, "_stream_id": "s1"}
|
||||
)
|
||||
await channel.send_delta(
|
||||
chat_id, "", {"_stream_end": True, "_stream_id": "s1"}
|
||||
)
|
||||
|
||||
delta1 = json.loads(await client.recv())
|
||||
assert delta1["event"] == "delta"
|
||||
assert delta1["text"] == "Hello "
|
||||
assert delta1["stream_id"] == "s1"
|
||||
|
||||
delta2 = json.loads(await client.recv())
|
||||
assert delta2["event"] == "delta"
|
||||
assert delta2["text"] == "world"
|
||||
assert delta2["stream_id"] == "s1"
|
||||
|
||||
end = json.loads(await client.recv())
|
||||
assert end["event"] == "stream_end"
|
||||
assert end["stream_id"] == "s1"
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_issue_rejects_when_at_capacity(bus: MagicMock) -> None:
|
||||
port = 29881
|
||||
channel = _ch(bus, port=port, tokenIssuePath="/auth/token", tokenIssueSecret="s")
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
# Fill issued tokens to capacity
|
||||
channel._issued_tokens = {
|
||||
f"nbwt_fill_{i}": time.monotonic() + 300 for i in range(channel._MAX_ISSUED_TOKENS)
|
||||
}
|
||||
|
||||
resp = await _http_get(
|
||||
f"http://127.0.0.1:{port}/auth/token",
|
||||
headers={"Authorization": "Bearer s"},
|
||||
)
|
||||
assert resp.status_code == 429
|
||||
data = resp.json()
|
||||
assert "error" in data
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allow_from_rejects_unauthorized_client_id(bus: MagicMock) -> None:
|
||||
port = 29882
|
||||
channel = _ch(bus, port=port, allowFrom=["alice", "bob"])
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=eve"):
|
||||
pass
|
||||
assert exc_info.value.response.status_code == 403
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_id_truncation(bus: MagicMock) -> None:
|
||||
port = 29883
|
||||
channel = _ch(bus, port=port)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
long_id = "x" * 200
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id={long_id}") as client:
|
||||
ready = json.loads(await client.recv())
|
||||
assert ready["client_id"] == "x" * 128
|
||||
assert len(ready["client_id"]) == 128
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_utf8_binary_frame_ignored(bus: MagicMock) -> None:
|
||||
port = 29884
|
||||
channel = _ch(bus, port=port)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=bin-test") as client:
|
||||
await client.recv() # consume ready
|
||||
# Send non-UTF-8 bytes
|
||||
await client.send(b"\xff\xfe\xfd")
|
||||
await asyncio.sleep(0.05)
|
||||
# publish_inbound should NOT have been called
|
||||
bus.publish_inbound.assert_not_awaited()
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_static_token_accepts_issued_token_as_fallback(bus: MagicMock) -> None:
|
||||
port = 29885
|
||||
channel = _ch(
|
||||
bus, port=port,
|
||||
token="static-secret",
|
||||
tokenIssuePath="/auth/token",
|
||||
tokenIssueSecret="route-secret",
|
||||
)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
# Get an issued token
|
||||
resp = await _http_get(
|
||||
f"http://127.0.0.1:{port}/auth/token",
|
||||
headers={"Authorization": "Bearer route-secret"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
issued_token = resp.json()["token"]
|
||||
|
||||
# Connect using issued token (not the static one)
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?token={issued_token}&client_id=caller") as client:
|
||||
ready = json.loads(await client.recv())
|
||||
assert ready["event"] == "ready"
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allow_from_empty_list_denies_all(bus: MagicMock) -> None:
|
||||
port = 29886
|
||||
channel = _ch(bus, port=port, allowFrom=[])
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=anyone"):
|
||||
pass
|
||||
assert exc_info.value.response.status_code == 403
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_requires_token_without_issue_path(bus: MagicMock) -> None:
|
||||
"""When websocket_requires_token is True but no token or issue path configured, all connections are rejected."""
|
||||
port = 29887
|
||||
channel = _ch(bus, port=port, websocketRequiresToken=True)
|
||||
|
||||
server_task = asyncio.create_task(channel.start())
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
try:
|
||||
# No token at all → 401
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=u"):
|
||||
pass
|
||||
assert exc_info.value.response.status_code == 401
|
||||
|
||||
# Wrong token → 401
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
|
||||
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=u&token=wrong"):
|
||||
pass
|
||||
assert exc_info.value.response.status_code == 401
|
||||
finally:
|
||||
await channel.stop()
|
||||
await server_task
|
||||
@@ -0,0 +1,478 @@
|
||||
"""Integration tests for the WebSocket channel using WsTestClient.
|
||||
|
||||
Complements the unit/lightweight tests in test_websocket_channel.py by covering
|
||||
multi-client scenarios, edge cases, and realistic usage patterns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import websockets
|
||||
|
||||
from nanobot.channels.websocket import WebSocketChannel
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from ws_test_client import WsTestClient, issue_token, issue_token_ok
|
||||
|
||||
|
||||
def _ch(bus: Any, port: int, **kw: Any) -> WebSocketChannel:
|
||||
cfg: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"allowFrom": ["*"],
|
||||
"host": "127.0.0.1",
|
||||
"port": port,
|
||||
"path": "/",
|
||||
"websocketRequiresToken": False,
|
||||
}
|
||||
cfg.update(kw)
|
||||
return WebSocketChannel(cfg, bus)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def bus() -> MagicMock:
|
||||
b = MagicMock()
|
||||
b.publish_inbound = AsyncMock()
|
||||
return b
|
||||
|
||||
|
||||
# -- Connection basics ----------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ready_event_fields(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29901)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29901/", client_id="c1") as c:
|
||||
r = await c.recv_ready()
|
||||
assert r.event == "ready"
|
||||
assert len(r.chat_id) == 36
|
||||
assert r.client_id == "c1"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anonymous_client_gets_generated_id(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29902)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29902/", client_id="") as c:
|
||||
r = await c.recv_ready()
|
||||
assert r.client_id.startswith("anon-")
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_connection_unique_chat_id(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29903)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29903/", client_id="a") as c1:
|
||||
async with WsTestClient("ws://127.0.0.1:29903/", client_id="b") as c2:
|
||||
assert (await c1.recv_ready()).chat_id != (await c2.recv_ready()).chat_id
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Inbound messages (client -> server) ----------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plain_text(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29904)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29904/", client_id="p") as c:
|
||||
await c.recv_ready()
|
||||
await c.send_text("hello world")
|
||||
await asyncio.sleep(0.1)
|
||||
inbound = bus.publish_inbound.call_args[0][0]
|
||||
assert inbound.content == "hello world"
|
||||
assert inbound.sender_id == "p"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_content_field(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29905)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29905/", client_id="j") as c:
|
||||
await c.recv_ready()
|
||||
await c.send_json({"content": "structured"})
|
||||
await asyncio.sleep(0.1)
|
||||
assert bus.publish_inbound.call_args[0][0].content == "structured"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_text_and_message_fields(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29906)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29906/", client_id="x") as c:
|
||||
await c.recv_ready()
|
||||
await c.send_json({"text": "via text"})
|
||||
await asyncio.sleep(0.1)
|
||||
assert bus.publish_inbound.call_args[0][0].content == "via text"
|
||||
await c.send_json({"message": "via message"})
|
||||
await asyncio.sleep(0.1)
|
||||
assert bus.publish_inbound.call_args[0][0].content == "via message"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_payload_ignored(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29907)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29907/", client_id="e") as c:
|
||||
await c.recv_ready()
|
||||
await c.send_text(" ")
|
||||
await c.send_json({})
|
||||
await asyncio.sleep(0.1)
|
||||
bus.publish_inbound.assert_not_awaited()
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_messages_preserve_order(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29908)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29908/", client_id="o") as c:
|
||||
await c.recv_ready()
|
||||
for i in range(5):
|
||||
await c.send_text(f"msg-{i}")
|
||||
await asyncio.sleep(0.2)
|
||||
contents = [call[0][0].content for call in bus.publish_inbound.call_args_list]
|
||||
assert contents == [f"msg-{i}" for i in range(5)]
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Outbound messages (server -> client) ---------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_send_message(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29909)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29909/", client_id="r") as c:
|
||||
ready = await c.recv_ready()
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=ready.chat_id, content="reply",
|
||||
))
|
||||
msg = await c.recv_message()
|
||||
assert msg.text == "reply"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_send_with_media_and_reply(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29910)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29910/", client_id="m") as c:
|
||||
ready = await c.recv_ready()
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=ready.chat_id, content="img",
|
||||
media=["/tmp/a.png"], reply_to="m1",
|
||||
))
|
||||
msg = await c.recv_message()
|
||||
assert msg.text == "img"
|
||||
assert msg.media == ["/tmp/a.png"]
|
||||
assert msg.reply_to == "m1"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Streaming ------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_deltas_and_end(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29911, streaming=True)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29911/", client_id="s") as c:
|
||||
cid = (await c.recv_ready()).chat_id
|
||||
for part in ("Hello", " ", "world", "!"):
|
||||
await ch.send_delta(cid, part, {"_stream_delta": True, "_stream_id": "s1"})
|
||||
await ch.send_delta(cid, "", {"_stream_end": True, "_stream_id": "s1"})
|
||||
|
||||
msgs = await c.collect_stream()
|
||||
deltas = [m for m in msgs if m.event == "delta"]
|
||||
assert "".join(d.text for d in deltas) == "Hello world!"
|
||||
ends = [m for m in msgs if m.event == "stream_end"]
|
||||
assert len(ends) == 1
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interleaved_streams(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29912, streaming=True)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29912/", client_id="i") as c:
|
||||
cid = (await c.recv_ready()).chat_id
|
||||
await ch.send_delta(cid, "A1", {"_stream_delta": True, "_stream_id": "sa"})
|
||||
await ch.send_delta(cid, "B1", {"_stream_delta": True, "_stream_id": "sb"})
|
||||
await ch.send_delta(cid, "A2", {"_stream_delta": True, "_stream_id": "sa"})
|
||||
await ch.send_delta(cid, "", {"_stream_end": True, "_stream_id": "sa"})
|
||||
await ch.send_delta(cid, "B2", {"_stream_delta": True, "_stream_id": "sb"})
|
||||
await ch.send_delta(cid, "", {"_stream_end": True, "_stream_id": "sb"})
|
||||
|
||||
msgs = await c.recv_n(6)
|
||||
sa = "".join(m.text for m in msgs if m.event == "delta" and m.stream_id == "sa")
|
||||
sb = "".join(m.text for m in msgs if m.event == "delta" and m.stream_id == "sb")
|
||||
assert sa == "A1A2"
|
||||
assert sb == "B1B2"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Multi-client ---------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_independent_sessions(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29913)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29913/", client_id="u1") as c1:
|
||||
async with WsTestClient("ws://127.0.0.1:29913/", client_id="u2") as c2:
|
||||
r1, r2 = await c1.recv_ready(), await c2.recv_ready()
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=r1.chat_id, content="for-u1",
|
||||
))
|
||||
assert (await c1.recv_message()).text == "for-u1"
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=r2.chat_id, content="for-u2",
|
||||
))
|
||||
assert (await c2.recv_message()).text == "for-u2"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnected_client_cleanup(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29914)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29914/", client_id="tmp") as c:
|
||||
chat_id = (await c.recv_ready()).chat_id
|
||||
# disconnected
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=chat_id, content="orphan",
|
||||
))
|
||||
assert chat_id not in ch._connections
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Authentication -------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_static_token_accepted(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29915, token="secret")
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29915/", client_id="a", token="secret") as c:
|
||||
assert (await c.recv_ready()).client_id == "a"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_static_token_rejected(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29916, token="correct")
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
|
||||
async with WsTestClient("ws://127.0.0.1:29916/", client_id="b", token="wrong"):
|
||||
pass
|
||||
assert exc.value.response.status_code == 401
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_issue_full_flow(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29917, path="/ws",
|
||||
tokenIssuePath="/auth/token", tokenIssueSecret="s",
|
||||
websocketRequiresToken=True)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
# no secret -> 401
|
||||
_, status = await issue_token(port=29917, issue_path="/auth/token")
|
||||
assert status == 401
|
||||
|
||||
# with secret -> token
|
||||
token = await issue_token_ok(port=29917, issue_path="/auth/token", secret="s")
|
||||
|
||||
# no token -> 401
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
|
||||
async with WsTestClient("ws://127.0.0.1:29917/ws", client_id="x"):
|
||||
pass
|
||||
assert exc.value.response.status_code == 401
|
||||
|
||||
# valid token -> ok
|
||||
async with WsTestClient("ws://127.0.0.1:29917/ws", client_id="ok", token=token) as c:
|
||||
assert (await c.recv_ready()).client_id == "ok"
|
||||
|
||||
# reuse -> 401
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
|
||||
async with WsTestClient("ws://127.0.0.1:29917/ws", client_id="r", token=token):
|
||||
pass
|
||||
assert exc.value.response.status_code == 401
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Path routing ---------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_path(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29918, path="/my-chat")
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29918/my-chat", client_id="p") as c:
|
||||
assert (await c.recv_ready()).event == "ready"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wrong_path_404(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29919, path="/ws")
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
|
||||
async with WsTestClient("ws://127.0.0.1:29919/wrong", client_id="x"):
|
||||
pass
|
||||
assert exc.value.response.status_code == 404
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trailing_slash_normalized(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29920, path="/ws")
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29920/ws/", client_id="s") as c:
|
||||
assert (await c.recv_ready()).event == "ready"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
# -- Edge cases -----------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_message(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29921)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29921/", client_id="big") as c:
|
||||
await c.recv_ready()
|
||||
big = "x" * 100_000
|
||||
await c.send_text(big)
|
||||
await asyncio.sleep(0.2)
|
||||
assert bus.publish_inbound.call_args[0][0].content == big
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unicode_roundtrip(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29922)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29922/", client_id="u") as c:
|
||||
ready = await c.recv_ready()
|
||||
text = "你好世界 🌍 日本語テスト"
|
||||
await c.send_text(text)
|
||||
await asyncio.sleep(0.1)
|
||||
assert bus.publish_inbound.call_args[0][0].content == text
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=ready.chat_id, content=text,
|
||||
))
|
||||
assert (await c.recv_message()).text == text
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rapid_fire(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29923)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29923/", client_id="r") as c:
|
||||
ready = await c.recv_ready()
|
||||
for i in range(50):
|
||||
await c.send_text(f"in-{i}")
|
||||
await asyncio.sleep(0.5)
|
||||
assert bus.publish_inbound.await_count == 50
|
||||
for i in range(50):
|
||||
await ch.send(OutboundMessage(
|
||||
channel="websocket", chat_id=ready.chat_id, content=f"out-{i}",
|
||||
))
|
||||
received = [(await c.recv_message()).text for _ in range(50)]
|
||||
assert received == [f"out-{i}" for i in range(50)]
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_json_as_plain_text(bus: MagicMock) -> None:
|
||||
ch = _ch(bus, 29924)
|
||||
t = asyncio.create_task(ch.start())
|
||||
await asyncio.sleep(0.3)
|
||||
try:
|
||||
async with WsTestClient("ws://127.0.0.1:29924/", client_id="j") as c:
|
||||
await c.recv_ready()
|
||||
await c.send_text("{broken json")
|
||||
await asyncio.sleep(0.1)
|
||||
assert bus.publish_inbound.call_args[0][0].content == "{broken json"
|
||||
finally:
|
||||
await ch.stop(); await t
|
||||
@@ -0,0 +1,584 @@
|
||||
"""Tests for WeCom channel: helpers, download, upload, send, and message processing."""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
import importlib.util
|
||||
|
||||
WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None
|
||||
except ImportError:
|
||||
WECOM_AVAILABLE = False
|
||||
|
||||
if not WECOM_AVAILABLE:
|
||||
pytest.skip("WeCom dependencies not installed (wecom_aibot_sdk)", allow_module_level=True)
|
||||
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.wecom import (
|
||||
WecomChannel,
|
||||
WecomConfig,
|
||||
_guess_wecom_media_type,
|
||||
_sanitize_filename,
|
||||
)
|
||||
|
||||
# Try to import the real response class; fall back to a stub if unavailable.
|
||||
try:
|
||||
from wecom_aibot_sdk.utils import WsResponse
|
||||
|
||||
_RealWsResponse = WsResponse
|
||||
except ImportError:
|
||||
_RealWsResponse = None
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
"""Minimal stand-in for wecom_aibot_sdk WsResponse."""
|
||||
|
||||
def __init__(self, errcode: int = 0, body: dict | None = None, errmsg: str = "ok"):
|
||||
self.errcode = errcode
|
||||
self.errmsg = errmsg
|
||||
self.body = body or {}
|
||||
|
||||
|
||||
class _FakeWsManager:
|
||||
"""Tracks send_reply calls and returns configurable responses."""
|
||||
|
||||
def __init__(self, responses: list[_FakeResponse] | None = None):
|
||||
self.responses = responses or []
|
||||
self.calls: list[tuple[str, dict, str]] = []
|
||||
self._idx = 0
|
||||
|
||||
async def send_reply(self, req_id: str, data: dict, cmd: str) -> _FakeResponse:
|
||||
self.calls.append((req_id, data, cmd))
|
||||
if self._idx < len(self.responses):
|
||||
resp = self.responses[self._idx]
|
||||
self._idx += 1
|
||||
return resp
|
||||
return _FakeResponse()
|
||||
|
||||
|
||||
class _FakeFrame:
|
||||
"""Minimal frame object with a body dict."""
|
||||
|
||||
def __init__(self, body: dict | None = None):
|
||||
self.body = body or {}
|
||||
|
||||
|
||||
class _FakeWeComClient:
|
||||
"""Fake WeCom client with mock methods."""
|
||||
|
||||
def __init__(self, ws_responses: list[_FakeResponse] | None = None):
|
||||
self._ws_manager = _FakeWsManager(ws_responses)
|
||||
self.download_file = AsyncMock(return_value=(None, None))
|
||||
self.reply = AsyncMock()
|
||||
self.reply_stream = AsyncMock()
|
||||
self.send_message = AsyncMock()
|
||||
self.reply_welcome = AsyncMock()
|
||||
|
||||
|
||||
# ── Helper function tests (pure, no async) ──────────────────────────
|
||||
|
||||
|
||||
def test_sanitize_filename_strips_path_traversal() -> None:
|
||||
assert _sanitize_filename("../../etc/passwd") == "passwd"
|
||||
|
||||
|
||||
def test_sanitize_filename_keeps_chinese_chars() -> None:
|
||||
assert _sanitize_filename("文件（1）.jpg") == "文件（1）.jpg"
|
||||
|
||||
|
||||
def test_sanitize_filename_empty_input() -> None:
|
||||
assert _sanitize_filename("") == ""
|
||||
|
||||
|
||||
def test_guess_wecom_media_type_image() -> None:
|
||||
for ext in (".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"):
|
||||
assert _guess_wecom_media_type(f"photo{ext}") == "image"
|
||||
|
||||
|
||||
def test_guess_wecom_media_type_video() -> None:
|
||||
for ext in (".mp4", ".avi", ".mov"):
|
||||
assert _guess_wecom_media_type(f"video{ext}") == "video"
|
||||
|
||||
|
||||
def test_guess_wecom_media_type_voice() -> None:
|
||||
for ext in (".amr", ".mp3", ".wav", ".ogg"):
|
||||
assert _guess_wecom_media_type(f"audio{ext}") == "voice"
|
||||
|
||||
|
||||
def test_guess_wecom_media_type_file_fallback() -> None:
|
||||
for ext in (".pdf", ".doc", ".xlsx", ".zip"):
|
||||
assert _guess_wecom_media_type(f"doc{ext}") == "file"
|
||||
|
||||
|
||||
def test_guess_wecom_media_type_case_insensitive() -> None:
|
||||
assert _guess_wecom_media_type("photo.PNG") == "image"
|
||||
assert _guess_wecom_media_type("photo.Jpg") == "image"
|
||||
|
||||
|
||||
# ── _download_and_save_media() ──────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_and_save_success() -> None:
|
||||
"""Successful download writes file and returns sanitized path."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
fake_data = b"\x89PNG\r\nfake image"
|
||||
client.download_file.return_value = (fake_data, "raw_photo.png")
|
||||
|
||||
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(tempfile.gettempdir())):
|
||||
path = await channel._download_and_save_media("https://example.com/img.png", "aes_key", "image", "photo.png")
|
||||
|
||||
assert path is not None
|
||||
assert os.path.isfile(path)
|
||||
assert os.path.basename(path) == "photo.png"
|
||||
# Cleanup
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_and_save_oversized_rejected() -> None:
|
||||
"""Data exceeding 200MB is rejected → returns None."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
big_data = b"\x00" * (200 * 1024 * 1024 + 1) # 200MB + 1 byte
|
||||
client.download_file.return_value = (big_data, "big.bin")
|
||||
|
||||
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(tempfile.gettempdir())):
|
||||
result = await channel._download_and_save_media("https://example.com/big.bin", "key", "file", "big.bin")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_and_save_failure() -> None:
|
||||
"""SDK returns None data → returns None."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
client.download_file.return_value = (None, None)
|
||||
|
||||
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(tempfile.gettempdir())):
|
||||
result = await channel._download_and_save_media("https://example.com/fail.png", "key", "image")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
# ── _upload_media_ws() ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_media_ws_success() -> None:
|
||||
"""Happy path: init → chunk → finish → returns (media_id, media_type)."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
tmp = f.name
|
||||
|
||||
try:
|
||||
responses = [
|
||||
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
|
||||
_FakeResponse(errcode=0, body={}),
|
||||
_FakeResponse(errcode=0, body={"media_id": "media_abc"}),
|
||||
]
|
||||
|
||||
client = _FakeWeComClient(responses)
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
channel._client = client
|
||||
|
||||
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
|
||||
media_id, media_type = await channel._upload_media_ws(client, tmp)
|
||||
|
||||
assert media_id == "media_abc"
|
||||
assert media_type == "image"
|
||||
finally:
|
||||
os.unlink(tmp)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_media_ws_oversized_file() -> None:
|
||||
"""File >200MB triggers ValueError → returns (None, None)."""
|
||||
# Instead of creating a real 200MB+ file, mock os.path.getsize and open
|
||||
with patch("os.path.getsize", return_value=200 * 1024 * 1024 + 1), \
|
||||
patch("builtins.open", MagicMock()):
|
||||
client = _FakeWeComClient()
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
channel._client = client
|
||||
|
||||
result = await channel._upload_media_ws(client, "/fake/large.bin")
|
||||
assert result == (None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_media_ws_init_failure() -> None:
|
||||
"""Init step returns errcode != 0 → returns (None, None)."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
|
||||
f.write(b"hello")
|
||||
tmp = f.name
|
||||
|
||||
try:
|
||||
responses = [
|
||||
_FakeResponse(errcode=50001, errmsg="invalid"),
|
||||
]
|
||||
|
||||
client = _FakeWeComClient(responses)
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
channel._client = client
|
||||
|
||||
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
|
||||
result = await channel._upload_media_ws(client, tmp)
|
||||
|
||||
assert result == (None, None)
|
||||
finally:
|
||||
os.unlink(tmp)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_media_ws_chunk_failure() -> None:
|
||||
"""Chunk step returns errcode != 0 → returns (None, None)."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
tmp = f.name
|
||||
|
||||
try:
|
||||
responses = [
|
||||
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
|
||||
_FakeResponse(errcode=50002, errmsg="chunk fail"),
|
||||
]
|
||||
|
||||
client = _FakeWeComClient(responses)
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
channel._client = client
|
||||
|
||||
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
|
||||
result = await channel._upload_media_ws(client, tmp)
|
||||
|
||||
assert result == (None, None)
|
||||
finally:
|
||||
os.unlink(tmp)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_media_ws_finish_no_media_id() -> None:
|
||||
"""Finish step returns empty media_id → returns (None, None)."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
tmp = f.name
|
||||
|
||||
try:
|
||||
responses = [
|
||||
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
|
||||
_FakeResponse(errcode=0, body={}),
|
||||
_FakeResponse(errcode=0, body={}), # no media_id
|
||||
]
|
||||
|
||||
client = _FakeWeComClient(responses)
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
channel._client = client
|
||||
|
||||
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
|
||||
result = await channel._upload_media_ws(client, tmp)
|
||||
|
||||
assert result == (None, None)
|
||||
finally:
|
||||
os.unlink(tmp)
|
||||
|
||||
|
||||
# ── send() ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_text_with_frame() -> None:
|
||||
"""When frame is stored, send uses reply_stream for final text."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
channel._generate_req_id = lambda x: f"req_{x}"
|
||||
channel._chat_frames["chat1"] = _FakeFrame()
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(channel="wecom", chat_id="chat1", content="hello")
|
||||
)
|
||||
|
||||
client.reply_stream.assert_called_once()
|
||||
call_args = client.reply_stream.call_args
|
||||
assert call_args[0][2] == "hello" # content arg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_progress_with_frame() -> None:
|
||||
"""When metadata has _progress, send uses reply_stream with finish=False."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
channel._generate_req_id = lambda x: f"req_{x}"
|
||||
channel._chat_frames["chat1"] = _FakeFrame()
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(channel="wecom", chat_id="chat1", content="thinking...", metadata={"_progress": True})
|
||||
)
|
||||
|
||||
client.reply_stream.assert_called_once()
|
||||
call_args = client.reply_stream.call_args
|
||||
assert call_args[0][2] == "thinking..." # content arg
|
||||
assert call_args[1]["finish"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_proactive_without_frame() -> None:
|
||||
"""Without stored frame, send uses send_message with markdown."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(channel="wecom", chat_id="chat1", content="proactive msg")
|
||||
)
|
||||
|
||||
client.send_message.assert_called_once()
|
||||
call_args = client.send_message.call_args
|
||||
assert call_args[0][0] == "chat1"
|
||||
assert call_args[0][1]["msgtype"] == "markdown"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_then_text() -> None:
|
||||
"""Media files are uploaded and sent before text content."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
tmp = f.name
|
||||
|
||||
try:
|
||||
responses = [
|
||||
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
|
||||
_FakeResponse(errcode=0, body={}),
|
||||
_FakeResponse(errcode=0, body={"media_id": "media_123"}),
|
||||
]
|
||||
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient(responses)
|
||||
channel._client = client
|
||||
channel._generate_req_id = lambda x: f"req_{x}"
|
||||
channel._chat_frames["chat1"] = _FakeFrame()
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(channel="wecom", chat_id="chat1", content="see image", media=[tmp])
|
||||
)
|
||||
|
||||
# Media should have been sent via reply
|
||||
media_calls = [c for c in client.reply.call_args_list if c[0][1].get("msgtype") == "image"]
|
||||
assert len(media_calls) == 1
|
||||
assert media_calls[0][0][1]["image"]["media_id"] == "media_123"
|
||||
|
||||
# Text should have been sent via reply_stream
|
||||
client.reply_stream.assert_called_once()
|
||||
finally:
|
||||
os.unlink(tmp)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_file_not_found() -> None:
|
||||
"""Non-existent media path is skipped with a warning."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
channel._generate_req_id = lambda x: f"req_{x}"
|
||||
channel._chat_frames["chat1"] = _FakeFrame()
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(channel="wecom", chat_id="chat1", content="hello", media=["/nonexistent/file.png"])
|
||||
)
|
||||
|
||||
# reply_stream should still be called for the text part
|
||||
client.reply_stream.assert_called_once()
|
||||
# No media reply should happen
|
||||
media_calls = [c for c in client.reply.call_args_list if c[0][1].get("msgtype") in ("image", "file", "video")]
|
||||
assert len(media_calls) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_exception_caught_not_raised() -> None:
|
||||
"""Exceptions inside send() must not propagate."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
channel._generate_req_id = lambda x: f"req_{x}"
|
||||
channel._chat_frames["chat1"] = _FakeFrame()
|
||||
|
||||
# Make reply_stream raise
|
||||
client.reply_stream.side_effect = RuntimeError("boom")
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(channel="wecom", chat_id="chat1", content="fail test")
|
||||
)
|
||||
# No exception — test passes if we reach here.
|
||||
|
||||
|
||||
# ── _process_message() ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_text_message() -> None:
|
||||
"""Text message is routed to bus with correct fields."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
frame = _FakeFrame(body={
|
||||
"msgid": "msg_text_1",
|
||||
"chatid": "chat1",
|
||||
"chattype": "single",
|
||||
"from": {"userid": "user1"},
|
||||
"text": {"content": "hello wecom"},
|
||||
})
|
||||
|
||||
await channel._process_message(frame, "text")
|
||||
|
||||
msg = await channel.bus.consume_inbound()
|
||||
assert msg.sender_id == "user1"
|
||||
assert msg.chat_id == "chat1"
|
||||
assert msg.content == "hello wecom"
|
||||
assert msg.metadata["msg_type"] == "text"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_image_message() -> None:
|
||||
"""Image message: download success → media_paths non-empty."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n")
|
||||
saved = f.name
|
||||
|
||||
client.download_file.return_value = (b"\x89PNG\r\n", "photo.png")
|
||||
channel._client = client
|
||||
|
||||
try:
|
||||
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(os.path.dirname(saved))):
|
||||
frame = _FakeFrame(body={
|
||||
"msgid": "msg_img_1",
|
||||
"chatid": "chat1",
|
||||
"from": {"userid": "user1"},
|
||||
"image": {"url": "https://example.com/img.png", "aeskey": "key123"},
|
||||
})
|
||||
await channel._process_message(frame, "image")
|
||||
|
||||
msg = await channel.bus.consume_inbound()
|
||||
assert len(msg.media) == 1
|
||||
assert msg.media[0].endswith("photo.png")
|
||||
assert "[image:" in msg.content
|
||||
finally:
|
||||
if os.path.exists(saved):
|
||||
pass # may have been overwritten; clean up if exists
|
||||
# Clean up any photo.png in tempdir
|
||||
p = os.path.join(os.path.dirname(saved), "photo.png")
|
||||
if os.path.exists(p):
|
||||
os.unlink(p)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_file_message() -> None:
|
||||
"""File message: download success → media_paths non-empty (critical fix verification)."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
|
||||
f.write(b"%PDF-1.4 fake")
|
||||
saved = f.name
|
||||
|
||||
client.download_file.return_value = (b"%PDF-1.4 fake", "report.pdf")
|
||||
channel._client = client
|
||||
|
||||
try:
|
||||
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(os.path.dirname(saved))):
|
||||
frame = _FakeFrame(body={
|
||||
"msgid": "msg_file_1",
|
||||
"chatid": "chat1",
|
||||
"from": {"userid": "user1"},
|
||||
"file": {"url": "https://example.com/report.pdf", "aeskey": "key456", "name": "report.pdf"},
|
||||
})
|
||||
await channel._process_message(frame, "file")
|
||||
|
||||
msg = await channel.bus.consume_inbound()
|
||||
assert len(msg.media) == 1
|
||||
assert msg.media[0].endswith("report.pdf")
|
||||
assert "[file: report.pdf]" in msg.content
|
||||
finally:
|
||||
p = os.path.join(os.path.dirname(saved), "report.pdf")
|
||||
if os.path.exists(p):
|
||||
os.unlink(p)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_voice_message() -> None:
|
||||
"""Voice message: transcribed text is included in content."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
frame = _FakeFrame(body={
|
||||
"msgid": "msg_voice_1",
|
||||
"chatid": "chat1",
|
||||
"from": {"userid": "user1"},
|
||||
"voice": {"content": "transcribed text here"},
|
||||
})
|
||||
|
||||
await channel._process_message(frame, "voice")
|
||||
|
||||
msg = await channel.bus.consume_inbound()
|
||||
assert "transcribed text here" in msg.content
|
||||
assert "[voice]" in msg.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_message_deduplication() -> None:
|
||||
"""Same msg_id is not processed twice."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
frame = _FakeFrame(body={
|
||||
"msgid": "msg_dup_1",
|
||||
"chatid": "chat1",
|
||||
"from": {"userid": "user1"},
|
||||
"text": {"content": "once"},
|
||||
})
|
||||
|
||||
await channel._process_message(frame, "text")
|
||||
await channel._process_message(frame, "text")
|
||||
|
||||
msg = await channel.bus.consume_inbound()
|
||||
assert msg.content == "once"
|
||||
|
||||
# Second message should not appear on the bus
|
||||
assert channel.bus.inbound.empty()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_message_empty_content_skipped() -> None:
|
||||
"""Message with empty content produces no bus message."""
|
||||
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
|
||||
client = _FakeWeComClient()
|
||||
channel._client = client
|
||||
|
||||
frame = _FakeFrame(body={
|
||||
"msgid": "msg_empty_1",
|
||||
"chatid": "chat1",
|
||||
"from": {"userid": "user1"},
|
||||
"text": {"content": ""},
|
||||
})
|
||||
|
||||
await channel._process_message(frame, "text")
|
||||
|
||||
assert channel.bus.inbound.empty()
|
||||
@@ -0,0 +1,227 @@
|
||||
"""Lightweight WebSocket test client for integration testing the nanobot WebSocket channel.
|
||||
|
||||
Provides an async ``WsTestClient`` class and token-issuance helpers that
|
||||
integration tests can import and use directly::
|
||||
|
||||
from ws_test_client import WsTestClient
|
||||
|
||||
async with WsTestClient("ws://127.0.0.1:8765/", client_id="t") as c:
|
||||
ready = await c.recv_ready()
|
||||
await c.send_text("hello")
|
||||
msg = await c.recv_message()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import websockets
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
||||
|
||||
@dataclass
|
||||
class WsMessage:
|
||||
"""A parsed message received from the WebSocket server."""
|
||||
|
||||
event: str
|
||||
raw: dict[str, Any] = field(repr=False)
|
||||
|
||||
@property
|
||||
def text(self) -> str | None:
|
||||
return self.raw.get("text")
|
||||
|
||||
@property
|
||||
def chat_id(self) -> str | None:
|
||||
return self.raw.get("chat_id")
|
||||
|
||||
@property
|
||||
def client_id(self) -> str | None:
|
||||
return self.raw.get("client_id")
|
||||
|
||||
@property
|
||||
def media(self) -> list[str] | None:
|
||||
return self.raw.get("media")
|
||||
|
||||
@property
|
||||
def reply_to(self) -> str | None:
|
||||
return self.raw.get("reply_to")
|
||||
|
||||
@property
|
||||
def stream_id(self) -> str | None:
|
||||
return self.raw.get("stream_id")
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, WsMessage):
|
||||
return NotImplemented
|
||||
return self.event == other.event and self.raw == other.raw
|
||||
|
||||
|
||||
class WsTestClient:
|
||||
"""Async WebSocket test client with helper methods for common operations.
|
||||
|
||||
Usage::
|
||||
|
||||
async with WsTestClient("ws://127.0.0.1:8765/", client_id="tester") as client:
|
||||
ready = await client.recv_ready()
|
||||
await client.send_text("hello")
|
||||
msg = await client.recv_message(timeout=5.0)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
uri: str,
|
||||
*,
|
||||
client_id: str = "test-client",
|
||||
token: str = "",
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
params: list[str] = []
|
||||
if client_id:
|
||||
params.append(f"client_id={client_id}")
|
||||
if token:
|
||||
params.append(f"token={token}")
|
||||
sep = "&" if "?" in uri else "?"
|
||||
self._uri = uri + sep + "&".join(params) if params else uri
|
||||
self._extra_headers = extra_headers
|
||||
self._ws: ClientConnection | None = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
self._ws = await websockets.connect(
|
||||
self._uri,
|
||||
additional_headers=self._extra_headers,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._ws:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
|
||||
async def __aenter__(self) -> WsTestClient:
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.close()
|
||||
|
||||
@property
|
||||
def ws(self) -> ClientConnection:
|
||||
assert self._ws is not None, "Client is not connected"
|
||||
return self._ws
|
||||
|
||||
# -- Receiving --------------------------------------------------------
|
||||
|
||||
async def recv_raw(self, timeout: float = 10.0) -> dict[str, Any]:
|
||||
"""Receive and parse one raw JSON message with timeout."""
|
||||
raw = await asyncio.wait_for(self.ws.recv(), timeout=timeout)
|
||||
return json.loads(raw)
|
||||
|
||||
async def recv(self, timeout: float = 10.0) -> WsMessage:
|
||||
"""Receive one message, returning a WsMessage wrapper."""
|
||||
data = await self.recv_raw(timeout)
|
||||
return WsMessage(event=data.get("event", ""), raw=data)
|
||||
|
||||
async def recv_ready(self, timeout: float = 5.0) -> WsMessage:
|
||||
"""Receive and validate the 'ready' event."""
|
||||
msg = await self.recv(timeout)
|
||||
assert msg.event == "ready", f"Expected 'ready' event, got '{msg.event}'"
|
||||
return msg
|
||||
|
||||
async def recv_message(self, timeout: float = 10.0) -> WsMessage:
|
||||
"""Receive and validate a 'message' event."""
|
||||
msg = await self.recv(timeout)
|
||||
assert msg.event == "message", f"Expected 'message' event, got '{msg.event}'"
|
||||
return msg
|
||||
|
||||
async def recv_delta(self, timeout: float = 10.0) -> WsMessage:
|
||||
"""Receive and validate a 'delta' event."""
|
||||
msg = await self.recv(timeout)
|
||||
assert msg.event == "delta", f"Expected 'delta' event, got '{msg.event}'"
|
||||
return msg
|
||||
|
||||
async def recv_stream_end(self, timeout: float = 10.0) -> WsMessage:
|
||||
"""Receive and validate a 'stream_end' event."""
|
||||
msg = await self.recv(timeout)
|
||||
assert msg.event == "stream_end", f"Expected 'stream_end' event, got '{msg.event}'"
|
||||
return msg
|
||||
|
||||
async def collect_stream(self, timeout: float = 10.0) -> list[WsMessage]:
|
||||
"""Collect all deltas and the final stream_end into a list."""
|
||||
messages: list[WsMessage] = []
|
||||
while True:
|
||||
msg = await self.recv(timeout)
|
||||
messages.append(msg)
|
||||
if msg.event == "stream_end":
|
||||
break
|
||||
return messages
|
||||
|
||||
async def recv_n(self, n: int, timeout: float = 10.0) -> list[WsMessage]:
|
||||
"""Receive exactly *n* messages."""
|
||||
return [await self.recv(timeout) for _ in range(n)]
|
||||
|
||||
# -- Sending ----------------------------------------------------------
|
||||
|
||||
async def send_text(self, text: str) -> None:
|
||||
"""Send a plain text frame."""
|
||||
await self.ws.send(text)
|
||||
|
||||
async def send_json(self, data: dict[str, Any]) -> None:
|
||||
"""Send a JSON frame."""
|
||||
await self.ws.send(json.dumps(data, ensure_ascii=False))
|
||||
|
||||
async def send_content(self, content: str) -> None:
|
||||
"""Send content in the preferred JSON format ``{"content": ...}``."""
|
||||
await self.send_json({"content": content})
|
||||
|
||||
# -- Connection introspection -----------------------------------------
|
||||
|
||||
@property
|
||||
def closed(self) -> bool:
|
||||
return self._ws is None or self._ws.closed
|
||||
|
||||
|
||||
# -- Token issuance helpers -----------------------------------------------
|
||||
|
||||
|
||||
async def issue_token(
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8765,
|
||||
issue_path: str = "/auth/token",
|
||||
secret: str = "",
|
||||
) -> tuple[dict[str, Any] | None, int]:
|
||||
"""Request a short-lived token from the token-issue HTTP endpoint.
|
||||
|
||||
Returns ``(parsed_json_or_None, status_code)``.
|
||||
"""
|
||||
url = f"http://{host}:{port}{issue_path}"
|
||||
headers: dict[str, str] = {}
|
||||
if secret:
|
||||
headers["Authorization"] = f"Bearer {secret}"
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
resp = await loop.run_in_executor(
|
||||
None, lambda: httpx.get(url, headers=headers, timeout=5.0)
|
||||
)
|
||||
try:
|
||||
data = resp.json()
|
||||
except Exception:
|
||||
data = None
|
||||
return data, resp.status_code
|
||||
|
||||
|
||||
async def issue_token_ok(
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8765,
|
||||
issue_path: str = "/auth/token",
|
||||
secret: str = "",
|
||||
) -> str:
|
||||
"""Request a token, asserting success, and return the token string."""
|
||||
(data, status) = await issue_token(host, port, issue_path, secret)
|
||||
assert status == 200, f"Token issue failed with status {status}"
|
||||
assert data is not None
|
||||
token = data["token"]
|
||||
assert token.startswith("nbwt_"), f"Unexpected token format: {token}"
|
||||
return token
|
||||
@@ -327,3 +327,240 @@ async def test_external_update_preserves_run_history_records(tmp_path):
|
||||
|
||||
fresh._running = True
|
||||
fresh._save_store()
|
||||
|
||||
|
||||
# ── timer race regression tests ──
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timer_execution_is_not_rolled_back_by_list_jobs_reload(tmp_path):
|
||||
"""list_jobs() during _on_timer should not replace the active store and re-run the same due job."""
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
calls: list[str] = []
|
||||
|
||||
async def on_job(job):
|
||||
calls.append(job.id)
|
||||
# Simulate frontend polling list_jobs while the timer callback is mid-execution.
|
||||
service.list_jobs(include_disabled=True)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
service = CronService(store_path, on_job=on_job)
|
||||
service._running = True
|
||||
service._load_store()
|
||||
service._arm_timer = lambda: None
|
||||
|
||||
job = service.add_job(
|
||||
name="race",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
)
|
||||
job.state.next_run_at_ms = max(1, int(time.time() * 1000) - 1_000)
|
||||
service._save_store()
|
||||
|
||||
await service._on_timer()
|
||||
await service._on_timer()
|
||||
|
||||
assert calls == [job.id]
|
||||
loaded = service.get_job(job.id)
|
||||
assert loaded is not None
|
||||
assert loaded.state.last_run_at_ms is not None
|
||||
assert loaded.state.next_run_at_ms is not None
|
||||
assert loaded.state.next_run_at_ms > loaded.state.last_run_at_ms
|
||||
|
||||
|
||||
# ── update_job tests ──
|
||||
|
||||
|
||||
def test_update_job_changes_name(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="old name",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
)
|
||||
result = service.update_job(job.id, name="new name")
|
||||
assert isinstance(result, CronJob)
|
||||
assert result.name == "new name"
|
||||
assert result.payload.message == "hello"
|
||||
|
||||
|
||||
def test_update_job_changes_schedule(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="sched",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
)
|
||||
old_next = job.state.next_run_at_ms
|
||||
|
||||
new_sched = CronSchedule(kind="every", every_ms=120_000)
|
||||
result = service.update_job(job.id, schedule=new_sched)
|
||||
assert isinstance(result, CronJob)
|
||||
assert result.schedule.every_ms == 120_000
|
||||
assert result.state.next_run_at_ms != old_next
|
||||
|
||||
|
||||
def test_update_job_changes_message(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="msg",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="old message",
|
||||
)
|
||||
result = service.update_job(job.id, message="new message")
|
||||
assert isinstance(result, CronJob)
|
||||
assert result.payload.message == "new message"
|
||||
|
||||
|
||||
def test_update_job_changes_cron_expression(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="cron-job",
|
||||
schedule=CronSchedule(kind="cron", expr="0 9 * * *", tz="UTC"),
|
||||
message="hello",
|
||||
)
|
||||
result = service.update_job(
|
||||
job.id,
|
||||
schedule=CronSchedule(kind="cron", expr="0 18 * * *", tz="UTC"),
|
||||
)
|
||||
assert isinstance(result, CronJob)
|
||||
assert result.schedule.expr == "0 18 * * *"
|
||||
assert result.state.next_run_at_ms is not None
|
||||
|
||||
|
||||
def test_update_job_not_found(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
result = service.update_job("nonexistent", name="x")
|
||||
assert result == "not_found"
|
||||
|
||||
|
||||
def test_update_job_rejects_system_job(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
service.register_system_job(CronJob(
|
||||
id="dream",
|
||||
name="dream",
|
||||
schedule=CronSchedule(kind="cron", expr="0 */2 * * *", tz="UTC"),
|
||||
payload=CronPayload(kind="system_event"),
|
||||
))
|
||||
result = service.update_job("dream", name="hacked")
|
||||
assert result == "protected"
|
||||
assert service.get_job("dream").name == "dream"
|
||||
|
||||
|
||||
def test_update_job_validates_schedule(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="validate",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
)
|
||||
with pytest.raises(ValueError, match="unknown timezone"):
|
||||
service.update_job(
|
||||
job.id,
|
||||
schedule=CronSchedule(kind="cron", expr="0 9 * * *", tz="Bad/Zone"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_job_preserves_run_history(tmp_path) -> None:
|
||||
import asyncio
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
service = CronService(store_path, on_job=lambda _: asyncio.sleep(0))
|
||||
job = service.add_job(
|
||||
name="hist",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
)
|
||||
await service.run_job(job.id)
|
||||
|
||||
result = service.update_job(job.id, name="renamed")
|
||||
assert isinstance(result, CronJob)
|
||||
assert len(result.state.run_history) == 1
|
||||
assert result.state.run_history[0].status == "ok"
|
||||
|
||||
|
||||
def test_update_job_offline_writes_action(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="offline",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
)
|
||||
service.update_job(job.id, name="updated-offline")
|
||||
|
||||
action_path = tmp_path / "cron" / "action.jsonl"
|
||||
assert action_path.exists()
|
||||
lines = [l for l in action_path.read_text().strip().split("\n") if l]
|
||||
last = json.loads(lines[-1])
|
||||
assert last["action"] == "update"
|
||||
assert last["params"]["name"] == "updated-offline"
|
||||
|
||||
|
||||
def test_update_job_sentinel_channel_and_to(tmp_path) -> None:
|
||||
"""Passing None clears channel/to; omitting leaves them unchanged."""
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
job = service.add_job(
|
||||
name="sentinel",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
channel="telegram",
|
||||
to="user123",
|
||||
)
|
||||
assert job.payload.channel == "telegram"
|
||||
assert job.payload.to == "user123"
|
||||
|
||||
result = service.update_job(job.id, name="renamed")
|
||||
assert isinstance(result, CronJob)
|
||||
assert result.payload.channel == "telegram"
|
||||
assert result.payload.to == "user123"
|
||||
|
||||
result = service.update_job(job.id, channel=None, to=None)
|
||||
assert isinstance(result, CronJob)
|
||||
assert result.payload.channel is None
|
||||
assert result.payload.to is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_jobs_during_on_job_does_not_cause_stale_reload(tmp_path) -> None:
|
||||
"""Regression: if the bot calls list_jobs (which reloads from disk) during
|
||||
on_job execution, the in-memory next_run_at_ms update must not be lost.
|
||||
Previously this caused an infinite re-trigger loop."""
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
execution_count = 0
|
||||
|
||||
async def on_job_that_lists(job):
|
||||
nonlocal execution_count
|
||||
execution_count += 1
|
||||
# Simulate the bot calling cron(action=list) mid-execution
|
||||
service.list_jobs()
|
||||
|
||||
service = CronService(store_path, on_job=on_job_that_lists, max_sleep_ms=100)
|
||||
await service.start()
|
||||
|
||||
# Add two jobs scheduled in the past so they're immediately due
|
||||
now_ms = int(time.time() * 1000)
|
||||
for name in ("job-a", "job-b"):
|
||||
service.add_job(
|
||||
name=name,
|
||||
schedule=CronSchedule(kind="every", every_ms=3_600_000),
|
||||
message="test",
|
||||
)
|
||||
# Force next_run to the past so _on_timer picks them up
|
||||
for job in service._store.jobs:
|
||||
job.state.next_run_at_ms = now_ms - 1000
|
||||
service._save_store()
|
||||
service._arm_timer()
|
||||
|
||||
# Let the timer fire once
|
||||
await asyncio.sleep(0.3)
|
||||
service.stop()
|
||||
|
||||
# Each job should have run exactly once, not looped
|
||||
assert execution_count == 2
|
||||
|
||||
# Verify next_run_at_ms was persisted correctly (in the future)
|
||||
raw = json.loads(store_path.read_text())
|
||||
for j in raw["jobs"]:
|
||||
next_run = j["state"]["nextRunAtMs"]
|
||||
assert next_run is not None
|
||||
assert next_run > now_ms, f"Job '{j['name']}' next_run should be in the future"
|
||||
|
||||
@@ -84,6 +84,34 @@ class TestEnforceRoleAlternation:
|
||||
tool_msgs = [m for m in result if m["role"] == "tool"]
|
||||
assert len(tool_msgs) == 2
|
||||
|
||||
def test_consecutive_assistant_keeps_later_tool_call_message(self):
|
||||
msgs = [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Previous reply"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
|
||||
{"role": "tool", "content": "result1", "tool_call_id": "1"},
|
||||
{"role": "user", "content": "Next"},
|
||||
]
|
||||
result = LLMProvider._enforce_role_alternation(msgs)
|
||||
assert result[1]["role"] == "assistant"
|
||||
assert result[1]["tool_calls"] == [{"id": "1"}]
|
||||
assert result[1]["content"] is None
|
||||
assert result[2]["role"] == "tool"
|
||||
|
||||
def test_consecutive_assistant_does_not_overwrite_existing_tool_call_message(self):
|
||||
msgs = [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
|
||||
{"role": "assistant", "content": "Later plain assistant"},
|
||||
{"role": "tool", "content": "result1", "tool_call_id": "1"},
|
||||
{"role": "user", "content": "Next"},
|
||||
]
|
||||
result = LLMProvider._enforce_role_alternation(msgs)
|
||||
assert result[1]["role"] == "assistant"
|
||||
assert result[1]["tool_calls"] == [{"id": "1"}]
|
||||
assert result[1]["content"] is None
|
||||
assert result[2]["role"] == "tool"
|
||||
|
||||
def test_non_string_content_uses_latest(self):
|
||||
msgs = [
|
||||
{"role": "user", "content": [{"type": "text", "text": "A"}]},
|
||||
|
||||
@@ -550,11 +550,40 @@ def test_openai_compat_preserves_message_level_reasoning_fields() -> None:
|
||||
{"role": "user", "content": "thanks"},
|
||||
])
|
||||
|
||||
assert sanitized[1]["content"] is None
|
||||
assert sanitized[1]["reasoning_content"] == "hidden"
|
||||
assert sanitized[1]["extra_content"] == {"debug": True}
|
||||
assert sanitized[1]["tool_calls"][0]["extra_content"] == {"google": {"thought_signature": "sig"}}
|
||||
|
||||
|
||||
def test_openai_compat_keeps_tool_calls_after_consecutive_assistant_messages() -> None:
|
||||
with patch("nanobot.providers.openai_compat_provider.AsyncOpenAI"):
|
||||
provider = OpenAICompatProvider()
|
||||
|
||||
sanitized = provider._sanitize_messages([
|
||||
{"role": "user", "content": "不错"},
|
||||
{"role": "assistant", "content": "对，破 4 万指日可待"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "<think>我再查一下</think>",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_function_akxp3wqzn7ph_1",
|
||||
"type": "function",
|
||||
"function": {"name": "exec", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_function_akxp3wqzn7ph_1", "name": "exec", "content": "ok"},
|
||||
{"role": "user", "content": "多少star了呢"},
|
||||
])
|
||||
|
||||
assert sanitized[1]["role"] == "assistant"
|
||||
assert sanitized[1]["content"] is None
|
||||
assert sanitized[1]["tool_calls"][0]["id"] == "3ec83c30d"
|
||||
assert sanitized[2]["tool_call_id"] == "3ec83c30d"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_compat_stream_watchdog_returns_error_on_stall(monkeypatch) -> None:
|
||||
monkeypatch.setenv("NANOBOT_STREAM_IDLE_TIMEOUT_S", "0")
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
import inspect
|
||||
from types import SimpleNamespace
|
||||
|
||||
|
||||
def test_sanitize_persisted_blocks_truncate_text_shadowing_regression() -> None:
|
||||
"""Regression: avoid bool param shadowing imported truncate_text.
|
||||
|
||||
Buggy behavior (historical):
|
||||
- loop.py imports `truncate_text` from helpers
|
||||
- `_sanitize_persisted_blocks(..., truncate_text: bool=...)` uses same name
|
||||
- when called with `truncate_text=True`, function body executes `truncate_text(text, ...)`
|
||||
which resolves to bool and raises `TypeError: 'bool' object is not callable`.
|
||||
|
||||
This test asserts the fixed API exists and truncation works without raising.
|
||||
"""
|
||||
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
|
||||
sig = inspect.signature(AgentLoop._sanitize_persisted_blocks)
|
||||
assert "should_truncate_text" in sig.parameters
|
||||
assert "truncate_text" not in sig.parameters
|
||||
|
||||
dummy = SimpleNamespace(max_tool_result_chars=5)
|
||||
content = [{"type": "text", "text": "0123456789"}]
|
||||
|
||||
out = AgentLoop._sanitize_persisted_blocks(dummy, content, should_truncate_text=True)
|
||||
assert isinstance(out, list)
|
||||
assert out and out[0]["type"] == "text"
|
||||
assert isinstance(out[0]["text"], str)
|
||||
assert out[0]["text"] != content[0]["text"]
|
||||
|
||||
@@ -0,0 +1,423 @@
|
||||
"""Tests for advanced EditFileTool enhancements inspired by claude-code:
|
||||
- Delete-line newline cleanup
|
||||
- Smart quote normalization (curly ↔ straight)
|
||||
- Quote style preservation in replacements
|
||||
- Indentation preservation when fallback match is trimmed
|
||||
- Trailing whitespace stripping for new_text
|
||||
- File size protection
|
||||
- Stale detection with content-equality fallback
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.filesystem import EditFileTool, ReadFileTool, _find_match
|
||||
from nanobot.agent.tools import file_state
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_file_state():
|
||||
file_state.clear()
|
||||
yield
|
||||
file_state.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Delete-line newline cleanup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeleteLineCleanup:
|
||||
"""When new_text='' and deleting a line, trailing newline should be consumed."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_line_consumes_trailing_newline(self, tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("line1\nline2\nline3\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="line2", new_text="")
|
||||
assert "Successfully" in result
|
||||
content = f.read_text()
|
||||
# Should not leave a blank line where line2 was
|
||||
assert content == "line1\nline3\n"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_line_with_explicit_newline_in_old_text(self, tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("line1\nline2\nline3\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="line2\n", new_text="")
|
||||
assert "Successfully" in result
|
||||
assert f.read_text() == "line1\nline3\n"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_preserves_content_when_not_trailing_newline(self, tool, tmp_path):
|
||||
"""Deleting a word mid-line should not consume extra characters."""
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world here\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="world ", new_text="")
|
||||
assert "Successfully" in result
|
||||
assert f.read_text() == "hello here\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Smart quote normalization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSmartQuoteNormalization:
|
||||
"""_find_match should handle curly ↔ straight quote fallback."""
|
||||
|
||||
def test_curly_double_quotes_match_straight(self):
|
||||
content = 'She said \u201chello\u201d to him'
|
||||
old_text = 'She said "hello" to him'
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match is not None
|
||||
assert count == 1
|
||||
# Returned match should be the ORIGINAL content with curly quotes
|
||||
assert "\u201c" in match
|
||||
|
||||
def test_curly_single_quotes_match_straight(self):
|
||||
content = "it\u2019s a test"
|
||||
old_text = "it's a test"
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match is not None
|
||||
assert count == 1
|
||||
assert "\u2019" in match
|
||||
|
||||
def test_straight_matches_curly_in_old_text(self):
|
||||
content = 'x = "hello"'
|
||||
old_text = 'x = \u201chello\u201d'
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match is not None
|
||||
assert count == 1
|
||||
|
||||
def test_exact_match_still_preferred_over_quote_normalization(self):
|
||||
content = 'x = "hello"'
|
||||
old_text = 'x = "hello"'
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match == old_text
|
||||
assert count == 1
|
||||
|
||||
|
||||
class TestQuoteStylePreservation:
|
||||
"""When quote-normalized matching occurs, replacement should preserve actual quote style."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replacement_preserves_curly_double_quotes(self, tool, tmp_path):
|
||||
f = tmp_path / "quotes.txt"
|
||||
f.write_text('message = “hello”\n', encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text='message = "hello"',
|
||||
new_text='message = "goodbye"',
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == 'message = “goodbye”\n'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replacement_preserves_curly_apostrophe(self, tool, tmp_path):
|
||||
f = tmp_path / "apostrophe.txt"
|
||||
f.write_text("it’s fine\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="it's fine",
|
||||
new_text="it's better",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == "it’s better\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Indentation preservation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIndentationPreservation:
|
||||
"""Replacement should keep outer indentation when trim fallback matched."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_fallback_preserves_outer_indentation(self, tool, tmp_path):
|
||||
f = tmp_path / "indent.py"
|
||||
f.write_text(
|
||||
"if True:\n"
|
||||
" def foo():\n"
|
||||
" pass\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="def foo():\n pass",
|
||||
new_text="def bar():\n return 1",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == (
|
||||
"if True:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Failure diagnostics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEditDiagnostics:
|
||||
"""Failure paths should offer actionable hints."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ambiguous_match_reports_candidate_lines(self, tool, tmp_path):
|
||||
f = tmp_path / "dup.py"
|
||||
f.write_text("aaa\nbbb\naaa\nbbb\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="aaa\nbbb", new_text="xxx")
|
||||
assert "appears 2 times" in result.lower()
|
||||
assert "line 1" in result.lower()
|
||||
assert "line 3" in result.lower()
|
||||
assert "replace_all=true" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_not_found_reports_whitespace_hint(self, tool, tmp_path):
|
||||
f = tmp_path / "space.py"
|
||||
f.write_text("value = 1\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="value = 1", new_text="value = 2")
|
||||
assert "Error" in result
|
||||
assert "whitespace" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_not_found_reports_case_hint(self, tool, tmp_path):
|
||||
f = tmp_path / "case.py"
|
||||
f.write_text("HelloWorld\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="helloworld", new_text="goodbye")
|
||||
assert "Error" in result
|
||||
assert "letter case differs" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Advanced fallback replacement behavior
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAdvancedReplaceAll:
|
||||
"""replace_all should work correctly for fallback-based matches too."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_all_preserves_each_match_indentation(self, tool, tmp_path):
|
||||
f = tmp_path / "indent_multi.py"
|
||||
f.write_text(
|
||||
"if a:\n"
|
||||
" def foo():\n"
|
||||
" pass\n"
|
||||
"if b:\n"
|
||||
" def foo():\n"
|
||||
" pass\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="def foo():\n pass",
|
||||
new_text="def bar():\n return 1",
|
||||
replace_all=True,
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == (
|
||||
"if a:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
"if b:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_and_quote_fallback_match_succeeds(self, tool, tmp_path):
|
||||
f = tmp_path / "quote_indent.py"
|
||||
f.write_text(" message = “hello”\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text='message = "hello"',
|
||||
new_text='message = "goodbye"',
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == " message = “goodbye”\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Advanced fallback replacement behavior
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAdvancedReplaceAll:
|
||||
"""replace_all should work correctly for fallback-based matches too."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_all_preserves_each_match_indentation(self, tool, tmp_path):
|
||||
f = tmp_path / "indent_multi.py"
|
||||
f.write_text(
|
||||
"if a:\n"
|
||||
" def foo():\n"
|
||||
" pass\n"
|
||||
"if b:\n"
|
||||
" def foo():\n"
|
||||
" pass\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="def foo():\n pass",
|
||||
new_text="def bar():\n return 1",
|
||||
replace_all=True,
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == (
|
||||
"if a:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
"if b:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_and_quote_fallback_match_succeeds(self, tool, tmp_path):
|
||||
f = tmp_path / "quote_indent.py"
|
||||
f.write_text(" message = “hello”\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text='message = "hello"',
|
||||
new_text='message = "goodbye"',
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == " message = “goodbye”\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trailing whitespace stripping on new_text
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTrailingWhitespaceStrip:
|
||||
"""new_text trailing whitespace should be stripped (except .md files)."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strips_trailing_whitespace_from_new_text(self, tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("x = 1\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f), old_text="x = 1", new_text="x = 2 \ny = 3 ",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
content = f.read_text()
|
||||
assert "x = 2\ny = 3\n" == content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preserves_trailing_whitespace_in_markdown(self, tool, tmp_path):
|
||||
f = tmp_path / "doc.md"
|
||||
f.write_text("# Title\n", encoding="utf-8")
|
||||
# Markdown uses trailing double-space for line breaks
|
||||
result = await tool.execute(
|
||||
path=str(f), old_text="# Title", new_text="# Title \nSubtitle ",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
content = f.read_text()
|
||||
# Trailing spaces should be preserved for markdown
|
||||
assert "Title " in content
|
||||
assert "Subtitle " in content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# File size protection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFileSizeProtection:
|
||||
"""Editing extremely large files should be rejected."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_file_over_size_limit(self, tool, tmp_path):
|
||||
f = tmp_path / "huge.txt"
|
||||
f.write_text("x", encoding="utf-8")
|
||||
# Monkey-patch the file size check by creating a stat mock
|
||||
original_stat = f.stat
|
||||
|
||||
class FakeStat:
|
||||
def __init__(self, real_stat):
|
||||
self._real = real_stat
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._real, name)
|
||||
|
||||
@property
|
||||
def st_size(self):
|
||||
return 2 * 1024 * 1024 * 1024 # 2 GiB
|
||||
|
||||
import unittest.mock
|
||||
with unittest.mock.patch.object(type(f), 'stat', return_value=FakeStat(f.stat())):
|
||||
result = await tool.execute(path=str(f), old_text="x", new_text="y")
|
||||
assert "Error" in result
|
||||
assert "too large" in result.lower() or "size" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stale detection with content-equality fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStaleDetectionContentFallback:
|
||||
"""When mtime changed but file content is unchanged, edit should proceed without warning."""
|
||||
|
||||
@pytest.fixture()
|
||||
def read_tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.fixture()
|
||||
def edit_tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mtime_bump_same_content_no_warning(self, read_tool, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
await read_tool.execute(path=str(f))
|
||||
|
||||
# Touch the file to bump mtime without changing content
|
||||
time.sleep(0.05)
|
||||
original_content = f.read_text()
|
||||
f.write_text(original_content, encoding="utf-8")
|
||||
|
||||
result = await edit_tool.execute(path=str(f), old_text="world", new_text="earth")
|
||||
assert "Successfully" in result
|
||||
# Should NOT warn about modification since content is the same
|
||||
assert "modified" not in result.lower()
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Tests for EditFileTool enhancements: read-before-edit tracking, path suggestions,
|
||||
.ipynb detection, and create-file semantics."""
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.filesystem import EditFileTool, ReadFileTool, WriteFileTool
|
||||
from nanobot.agent.tools import file_state
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_file_state():
|
||||
"""Reset global read-state between tests."""
|
||||
file_state.clear()
|
||||
yield
|
||||
file_state.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Read-before-edit tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditReadTracking:
|
||||
"""edit_file should warn when file hasn't been read first."""
|
||||
|
||||
@pytest.fixture()
|
||||
def read_tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.fixture()
|
||||
def edit_tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_warns_if_file_not_read_first(self, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
result = await edit_tool.execute(path=str(f), old_text="world", new_text="earth")
|
||||
# Should still succeed but include a warning
|
||||
assert "Successfully" in result
|
||||
assert "not been read" in result.lower() or "warning" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_succeeds_cleanly_after_read(self, read_tool, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
await read_tool.execute(path=str(f))
|
||||
result = await edit_tool.execute(path=str(f), old_text="world", new_text="earth")
|
||||
assert "Successfully" in result
|
||||
# No warning when file was read first
|
||||
assert "not been read" not in result.lower()
|
||||
assert f.read_text() == "hello earth"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_warns_if_file_modified_since_read(self, read_tool, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
await read_tool.execute(path=str(f))
|
||||
# External modification
|
||||
f.write_text("hello universe", encoding="utf-8")
|
||||
result = await edit_tool.execute(path=str(f), old_text="universe", new_text="earth")
|
||||
assert "Successfully" in result
|
||||
assert "modified" in result.lower() or "warning" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Create-file semantics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditCreateFile:
|
||||
"""edit_file with old_text='' creates new file if not exists."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_new_file_with_empty_old_text(self, tool, tmp_path):
|
||||
f = tmp_path / "subdir" / "new.py"
|
||||
result = await tool.execute(path=str(f), old_text="", new_text="print('hi')")
|
||||
assert "created" in result.lower() or "Successfully" in result
|
||||
assert f.exists()
|
||||
assert f.read_text() == "print('hi')"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_fails_if_file_already_exists_and_not_empty(self, tool, tmp_path):
|
||||
f = tmp_path / "existing.py"
|
||||
f.write_text("existing content", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="", new_text="new content")
|
||||
assert "Error" in result or "already exists" in result.lower()
|
||||
# File should be unchanged
|
||||
assert f.read_text() == "existing content"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_succeeds_if_file_exists_but_empty(self, tool, tmp_path):
|
||||
f = tmp_path / "empty.py"
|
||||
f.write_text("", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="", new_text="print('hi')")
|
||||
assert "Successfully" in result
|
||||
assert f.read_text() == "print('hi')"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# .ipynb detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditIpynbDetection:
|
||||
"""edit_file should refuse .ipynb and suggest notebook_edit."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ipynb_rejected_with_suggestion(self, tool, tmp_path):
|
||||
f = tmp_path / "analysis.ipynb"
|
||||
f.write_text('{"cells": []}', encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="x", new_text="y")
|
||||
assert "notebook" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path suggestion on not-found
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditPathSuggestion:
|
||||
"""edit_file should suggest similar paths on not-found."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_suggests_similar_filename(self, tool, tmp_path):
|
||||
f = tmp_path / "config.py"
|
||||
f.write_text("x = 1", encoding="utf-8")
|
||||
# Typo: conifg.py
|
||||
result = await tool.execute(
|
||||
path=str(tmp_path / "conifg.py"), old_text="x = 1", new_text="x = 2",
|
||||
)
|
||||
assert "Error" in result
|
||||
assert "config.py" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shows_cwd_in_error(self, tool, tmp_path):
|
||||
result = await tool.execute(
|
||||
path=str(tmp_path / "nonexistent.py"), old_text="a", new_text="b",
|
||||
)
|
||||
assert "Error" in result
|
||||
@@ -43,3 +43,34 @@ async def test_exec_path_append_preserves_system_path():
|
||||
tool = ExecTool(path_append="/opt/custom/bin")
|
||||
result = await tool.execute(command="ls /")
|
||||
assert "Exit code: 0" in result
|
||||
|
||||
|
||||
@_UNIX_ONLY
|
||||
@pytest.mark.asyncio
|
||||
async def test_exec_allowed_env_keys_passthrough(monkeypatch):
|
||||
"""Env vars listed in allowed_env_keys should be visible to commands."""
|
||||
monkeypatch.setenv("MY_CUSTOM_VAR", "hello-from-config")
|
||||
tool = ExecTool(allowed_env_keys=["MY_CUSTOM_VAR"])
|
||||
result = await tool.execute(command="printenv MY_CUSTOM_VAR")
|
||||
assert "hello-from-config" in result
|
||||
|
||||
|
||||
@_UNIX_ONLY
|
||||
@pytest.mark.asyncio
|
||||
async def test_exec_allowed_env_keys_does_not_leak_others(monkeypatch):
|
||||
"""Env vars NOT in allowed_env_keys should still be blocked."""
|
||||
monkeypatch.setenv("MY_CUSTOM_VAR", "hello-from-config")
|
||||
monkeypatch.setenv("MY_SECRET_VAR", "secret-value")
|
||||
tool = ExecTool(allowed_env_keys=["MY_CUSTOM_VAR"])
|
||||
result = await tool.execute(command="printenv MY_SECRET_VAR")
|
||||
assert "secret-value" not in result
|
||||
|
||||
|
||||
@_UNIX_ONLY
|
||||
@pytest.mark.asyncio
|
||||
async def test_exec_allowed_env_keys_missing_var_ignored(monkeypatch):
|
||||
"""If an allowed key is not set in the parent process, it should be silently skipped."""
|
||||
monkeypatch.delenv("NONEXISTENT_VAR_12345", raising=False)
|
||||
tool = ExecTool(allowed_env_keys=["NONEXISTENT_VAR_12345"])
|
||||
result = await tool.execute(command="printenv NONEXISTENT_VAR_12345")
|
||||
assert "Exit code: 1" in result
|
||||
|
||||
@@ -271,15 +271,11 @@ async def test_connect_mcp_servers_enabled_tools_supports_raw_names(
|
||||
) -> None:
|
||||
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
|
||||
registry = ToolRegistry()
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
try:
|
||||
await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=["demo"])},
|
||||
registry,
|
||||
stack,
|
||||
)
|
||||
finally:
|
||||
stacks = await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=["demo"])},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert registry.tool_names == ["mcp_test_demo"]
|
||||
@@ -291,15 +287,11 @@ async def test_connect_mcp_servers_enabled_tools_defaults_to_all(
|
||||
) -> None:
|
||||
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
|
||||
registry = ToolRegistry()
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
try:
|
||||
await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake")},
|
||||
registry,
|
||||
stack,
|
||||
)
|
||||
finally:
|
||||
stacks = await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake")},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert registry.tool_names == ["mcp_test_demo", "mcp_test_other"]
|
||||
@@ -311,15 +303,11 @@ async def test_connect_mcp_servers_enabled_tools_supports_wrapped_names(
|
||||
) -> None:
|
||||
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
|
||||
registry = ToolRegistry()
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
try:
|
||||
await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=["mcp_test_demo"])},
|
||||
registry,
|
||||
stack,
|
||||
)
|
||||
finally:
|
||||
stacks = await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=["mcp_test_demo"])},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert registry.tool_names == ["mcp_test_demo"]
|
||||
@@ -331,15 +319,11 @@ async def test_connect_mcp_servers_enabled_tools_empty_list_registers_none(
|
||||
) -> None:
|
||||
fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
|
||||
registry = ToolRegistry()
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
try:
|
||||
await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=[])},
|
||||
registry,
|
||||
stack,
|
||||
)
|
||||
finally:
|
||||
stacks = await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=[])},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert registry.tool_names == []
|
||||
@@ -358,15 +342,11 @@ async def test_connect_mcp_servers_enabled_tools_warns_on_unknown_entries(
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.tools.mcp.logger.warning", _warning)
|
||||
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
try:
|
||||
await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=["unknown"])},
|
||||
registry,
|
||||
stack,
|
||||
)
|
||||
finally:
|
||||
stacks = await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake", enabled_tools=["unknown"])},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert registry.tool_names == []
|
||||
@@ -376,6 +356,46 @@ async def test_connect_mcp_servers_enabled_tools_warns_on_unknown_entries(
|
||||
assert "Available wrapped names: mcp_test_demo" in warnings[-1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_mcp_servers_one_failure_does_not_block_others(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
sessions = {"good": _make_fake_session(["demo"])}
|
||||
|
||||
class _SelectiveClientSession:
|
||||
def __init__(self, read: object, _write: object) -> None:
|
||||
self._session = sessions[read]
|
||||
|
||||
async def __aenter__(self) -> object:
|
||||
return self._session
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> bool:
|
||||
return False
|
||||
|
||||
@asynccontextmanager
|
||||
async def _selective_stdio_client(params: object):
|
||||
if params.command == "bad":
|
||||
raise RuntimeError("boom")
|
||||
yield params.command, object()
|
||||
|
||||
monkeypatch.setattr(sys.modules["mcp"], "ClientSession", _SelectiveClientSession)
|
||||
monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _selective_stdio_client)
|
||||
|
||||
registry = ToolRegistry()
|
||||
stacks = await connect_mcp_servers(
|
||||
{
|
||||
"good": MCPServerConfig(command="good"),
|
||||
"bad": MCPServerConfig(command="bad"),
|
||||
},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert registry.tool_names == ["mcp_good_demo"]
|
||||
assert set(stacks) == {"good"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCPResourceWrapper tests
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -389,9 +409,7 @@ def _make_resource_def(
|
||||
return SimpleNamespace(name=name, uri=uri, description=description)
|
||||
|
||||
|
||||
def _make_resource_wrapper(
|
||||
session: object, *, timeout: float = 0.1
|
||||
) -> MCPResourceWrapper:
|
||||
def _make_resource_wrapper(session: object, *, timeout: float = 0.1) -> MCPResourceWrapper:
|
||||
return MCPResourceWrapper(session, "srv", _make_resource_def(), resource_timeout=timeout)
|
||||
|
||||
|
||||
@@ -434,9 +452,7 @@ async def test_resource_wrapper_execute_handles_timeout() -> None:
|
||||
await asyncio.sleep(1)
|
||||
return SimpleNamespace(contents=[])
|
||||
|
||||
wrapper = _make_resource_wrapper(
|
||||
SimpleNamespace(read_resource=read_resource), timeout=0.01
|
||||
)
|
||||
wrapper = _make_resource_wrapper(SimpleNamespace(read_resource=read_resource), timeout=0.01)
|
||||
result = await wrapper.execute()
|
||||
assert result == "(MCP resource read timed out after 0.01s)"
|
||||
|
||||
@@ -464,20 +480,14 @@ def _make_prompt_def(
|
||||
return SimpleNamespace(name=name, description=description, arguments=arguments)
|
||||
|
||||
|
||||
def _make_prompt_wrapper(
|
||||
session: object, *, timeout: float = 0.1
|
||||
) -> MCPPromptWrapper:
|
||||
return MCPPromptWrapper(
|
||||
session, "srv", _make_prompt_def(), prompt_timeout=timeout
|
||||
)
|
||||
def _make_prompt_wrapper(session: object, *, timeout: float = 0.1) -> MCPPromptWrapper:
|
||||
return MCPPromptWrapper(session, "srv", _make_prompt_def(), prompt_timeout=timeout)
|
||||
|
||||
|
||||
def test_prompt_wrapper_properties() -> None:
|
||||
arg1 = SimpleNamespace(name="topic", required=True)
|
||||
arg2 = SimpleNamespace(name="style", required=False)
|
||||
wrapper = MCPPromptWrapper(
|
||||
None, "myserver", _make_prompt_def(arguments=[arg1, arg2])
|
||||
)
|
||||
wrapper = MCPPromptWrapper(None, "myserver", _make_prompt_def(arguments=[arg1, arg2]))
|
||||
assert wrapper.name == "mcp_myserver_prompt_myprompt"
|
||||
assert "[MCP Prompt]" in wrapper.description
|
||||
assert "A test prompt" in wrapper.description
|
||||
@@ -528,9 +538,7 @@ async def test_prompt_wrapper_execute_handles_timeout() -> None:
|
||||
await asyncio.sleep(1)
|
||||
return SimpleNamespace(messages=[])
|
||||
|
||||
wrapper = _make_prompt_wrapper(
|
||||
SimpleNamespace(get_prompt=get_prompt), timeout=0.01
|
||||
)
|
||||
wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt), timeout=0.01)
|
||||
result = await wrapper.execute()
|
||||
assert result == "(MCP prompt call timed out after 0.01s)"
|
||||
|
||||
@@ -616,15 +624,11 @@ async def test_connect_registers_resources_and_prompts(
|
||||
prompt_names=["prompt_c"],
|
||||
)
|
||||
registry = ToolRegistry()
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
try:
|
||||
await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake")},
|
||||
registry,
|
||||
stack,
|
||||
)
|
||||
finally:
|
||||
stacks = await connect_mcp_servers(
|
||||
{"test": MCPServerConfig(command="fake")},
|
||||
registry,
|
||||
)
|
||||
for stack in stacks.values():
|
||||
await stack.aclose()
|
||||
|
||||
assert "mcp_test_tool_a" in registry.tool_names
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Test message tool suppress logic for final replies."""
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
@@ -86,6 +87,42 @@ class TestMessageToolSuppressLogic:
|
||||
assert result is not None
|
||||
assert "Hello" in result.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_injected_followup_with_message_tool_does_not_emit_empty_fallback(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
loop = _make_loop(tmp_path)
|
||||
tool_call = ToolCallRequest(
|
||||
id="call1", name="message",
|
||||
arguments={"content": "Tool reply", "channel": "feishu", "chat_id": "chat123"},
|
||||
)
|
||||
calls = iter([
|
||||
LLMResponse(content="First answer", tool_calls=[]),
|
||||
LLMResponse(content="", tool_calls=[tool_call]),
|
||||
LLMResponse(content="", tool_calls=[]),
|
||||
LLMResponse(content="", tool_calls=[]),
|
||||
LLMResponse(content="", tool_calls=[]),
|
||||
])
|
||||
loop.provider.chat_with_retry = AsyncMock(side_effect=lambda *a, **kw: next(calls))
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
|
||||
sent: list[OutboundMessage] = []
|
||||
mt = loop.tools.get("message")
|
||||
if isinstance(mt, MessageTool):
|
||||
mt.set_send_callback(AsyncMock(side_effect=lambda m: sent.append(m)))
|
||||
|
||||
pending_queue = asyncio.Queue()
|
||||
await pending_queue.put(
|
||||
InboundMessage(channel="feishu", sender_id="user1", chat_id="chat123", content="follow-up")
|
||||
)
|
||||
|
||||
msg = InboundMessage(channel="feishu", sender_id="user1", chat_id="chat123", content="Start")
|
||||
result = await loop._process_message(msg, pending_queue=pending_queue)
|
||||
|
||||
assert len(sent) == 1
|
||||
assert sent[0].content == "Tool reply"
|
||||
assert result is None
|
||||
|
||||
async def test_progress_hides_internal_reasoning(self, tmp_path: Path) -> None:
|
||||
loop = _make_loop(tmp_path)
|
||||
tool_call = ToolCallRequest(id="call1", name="read_file", arguments={"path": "foo.txt"})
|
||||
@@ -107,7 +144,7 @@ class TestMessageToolSuppressLogic:
|
||||
async def on_progress(content: str, *, tool_hint: bool = False) -> None:
|
||||
progress.append((content, tool_hint))
|
||||
|
||||
final_content, _, _ = await loop._run_agent_loop([], on_progress=on_progress)
|
||||
final_content, _, _, _, _ = await loop._run_agent_loop([], on_progress=on_progress)
|
||||
|
||||
assert final_content == "Done"
|
||||
assert progress == [
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Tests for NotebookEditTool — Jupyter .ipynb editing."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.notebook import NotebookEditTool
|
||||
|
||||
|
||||
def _make_notebook(cells: list[dict] | None = None, nbformat: int = 4, nbformat_minor: int = 5) -> dict:
|
||||
"""Build a minimal valid .ipynb structure."""
|
||||
return {
|
||||
"nbformat": nbformat,
|
||||
"nbformat_minor": nbformat_minor,
|
||||
"metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
|
||||
"cells": cells or [],
|
||||
}
|
||||
|
||||
|
||||
def _code_cell(source: str, cell_id: str | None = None) -> dict:
|
||||
cell = {"cell_type": "code", "source": source, "metadata": {}, "outputs": [], "execution_count": None}
|
||||
if cell_id:
|
||||
cell["id"] = cell_id
|
||||
return cell
|
||||
|
||||
|
||||
def _md_cell(source: str, cell_id: str | None = None) -> dict:
|
||||
cell = {"cell_type": "markdown", "source": source, "metadata": {}}
|
||||
if cell_id:
|
||||
cell["id"] = cell_id
|
||||
return cell
|
||||
|
||||
|
||||
def _write_nb(tmp_path, name: str, nb: dict) -> str:
|
||||
p = tmp_path / name
|
||||
p.write_text(json.dumps(nb), encoding="utf-8")
|
||||
return str(p)
|
||||
|
||||
|
||||
class TestNotebookEdit:
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return NotebookEditTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_cell_content(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("print('hello')"), _code_cell("x = 1")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="print('world')")
|
||||
assert "Successfully" in result
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert saved["cells"][0]["source"] == "print('world')"
|
||||
assert saved["cells"][1]["source"] == "x = 1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_cell_after_target(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("cell 0"), _code_cell("cell 1")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="inserted", edit_mode="insert")
|
||||
assert "Successfully" in result
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert len(saved["cells"]) == 3
|
||||
assert saved["cells"][0]["source"] == "cell 0"
|
||||
assert saved["cells"][1]["source"] == "inserted"
|
||||
assert saved["cells"][2]["source"] == "cell 1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cell(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("A"), _code_cell("B"), _code_cell("C")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=1, edit_mode="delete")
|
||||
assert "Successfully" in result
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert len(saved["cells"]) == 2
|
||||
assert saved["cells"][0]["source"] == "A"
|
||||
assert saved["cells"][1]["source"] == "C"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_new_notebook_from_scratch(self, tool, tmp_path):
|
||||
path = str(tmp_path / "new.ipynb")
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="# Hello", edit_mode="insert", cell_type="markdown")
|
||||
assert "Successfully" in result or "created" in result.lower()
|
||||
saved = json.loads((tmp_path / "new.ipynb").read_text())
|
||||
assert saved["nbformat"] == 4
|
||||
assert len(saved["cells"]) == 1
|
||||
assert saved["cells"][0]["cell_type"] == "markdown"
|
||||
assert saved["cells"][0]["source"] == "# Hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_cell_index_error(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("only cell")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=5, new_source="x")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_ipynb_rejected(self, tool, tmp_path):
|
||||
f = tmp_path / "script.py"
|
||||
f.write_text("pass")
|
||||
result = await tool.execute(path=str(f), cell_index=0, new_source="x")
|
||||
assert "Error" in result
|
||||
assert ".ipynb" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preserves_metadata_and_outputs(self, tool, tmp_path):
|
||||
cell = _code_cell("old")
|
||||
cell["outputs"] = [{"output_type": "stream", "text": "hello\n"}]
|
||||
cell["execution_count"] = 42
|
||||
nb = _make_notebook([cell])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
await tool.execute(path=path, cell_index=0, new_source="new")
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert saved["metadata"]["kernelspec"]["language"] == "python"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nbformat_45_generates_cell_id(self, tool, tmp_path):
|
||||
nb = _make_notebook([], nbformat_minor=5)
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
await tool.execute(path=path, cell_index=0, new_source="x = 1", edit_mode="insert")
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert "id" in saved["cells"][0]
|
||||
assert len(saved["cells"][0]["id"]) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_with_cell_type_markdown(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("code")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
await tool.execute(path=path, cell_index=0, new_source="# Title", edit_mode="insert", cell_type="markdown")
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert saved["cells"][1]["cell_type"] == "markdown"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_edit_mode_rejected(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("code")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="x", edit_mode="replcae")
|
||||
assert "Error" in result
|
||||
assert "edit_mode" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_cell_type_rejected(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("code")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="x", cell_type="raw")
|
||||
assert "Error" in result
|
||||
assert "cell_type" in result
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Tests for ReadFileTool enhancements: description fix, read dedup, PDF support, device blacklist."""
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.filesystem import ReadFileTool, WriteFileTool
|
||||
from nanobot.agent.tools import file_state
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_file_state():
|
||||
file_state.clear()
|
||||
yield
|
||||
file_state.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Description fix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadDescriptionFix:
|
||||
|
||||
def test_description_mentions_image_support(self):
|
||||
tool = ReadFileTool()
|
||||
assert "image" in tool.description.lower()
|
||||
|
||||
def test_description_no_longer_says_cannot_read_images(self):
|
||||
tool = ReadFileTool()
|
||||
assert "cannot read binary files or images" not in tool.description.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Read deduplication
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadDedup:
|
||||
"""Same file + same offset/limit + unchanged mtime -> short stub."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.fixture()
|
||||
def write_tool(self, tmp_path):
|
||||
return WriteFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_read_returns_unchanged_stub(self, tool, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("\n".join(f"line {i}" for i in range(100)), encoding="utf-8")
|
||||
first = await tool.execute(path=str(f))
|
||||
assert "line 0" in first
|
||||
second = await tool.execute(path=str(f))
|
||||
assert "unchanged" in second.lower()
|
||||
# Stub should not contain file content
|
||||
assert "line 0" not in second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_after_external_modification_returns_full(self, tool, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("original", encoding="utf-8")
|
||||
await tool.execute(path=str(f))
|
||||
# Modify the file externally
|
||||
f.write_text("modified content", encoding="utf-8")
|
||||
second = await tool.execute(path=str(f))
|
||||
assert "modified content" in second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_offset_returns_full(self, tool, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("\n".join(f"line {i}" for i in range(1, 21)), encoding="utf-8")
|
||||
await tool.execute(path=str(f), offset=1, limit=5)
|
||||
second = await tool.execute(path=str(f), offset=6, limit=5)
|
||||
# Different offset → full read, not stub
|
||||
assert "line 6" in second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_read_after_write_returns_full_content(self, tool, write_tool, tmp_path):
|
||||
f = tmp_path / "fresh.txt"
|
||||
result = await write_tool.execute(path=str(f), content="hello")
|
||||
assert "Successfully" in result
|
||||
read_result = await tool.execute(path=str(f))
|
||||
assert "hello" in read_result
|
||||
assert "unchanged" not in read_result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dedup_does_not_apply_to_images(self, tool, tmp_path):
|
||||
f = tmp_path / "img.png"
|
||||
f.write_bytes(b"\x89PNG\r\n\x1a\nfake-png-data")
|
||||
first = await tool.execute(path=str(f))
|
||||
assert isinstance(first, list)
|
||||
second = await tool.execute(path=str(f))
|
||||
# Images should always return full content blocks, not a stub
|
||||
assert isinstance(second, list)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDF support
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadPdf:
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pdf_returns_text_content(self, tool, tmp_path):
|
||||
fitz = pytest.importorskip("fitz")
|
||||
pdf_path = tmp_path / "test.pdf"
|
||||
doc = fitz.open()
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), "Hello PDF World")
|
||||
doc.save(str(pdf_path))
|
||||
doc.close()
|
||||
|
||||
result = await tool.execute(path=str(pdf_path))
|
||||
assert "Hello PDF World" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pdf_pages_parameter(self, tool, tmp_path):
|
||||
fitz = pytest.importorskip("fitz")
|
||||
pdf_path = tmp_path / "multi.pdf"
|
||||
doc = fitz.open()
|
||||
for i in range(5):
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), f"Page {i + 1} content")
|
||||
doc.save(str(pdf_path))
|
||||
doc.close()
|
||||
|
||||
result = await tool.execute(path=str(pdf_path), pages="2-3")
|
||||
assert "Page 2 content" in result
|
||||
assert "Page 3 content" in result
|
||||
assert "Page 1 content" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pdf_file_not_found_error(self, tool, tmp_path):
|
||||
result = await tool.execute(path=str(tmp_path / "nope.pdf"))
|
||||
assert "Error" in result
|
||||
assert "not found" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Device path blacklist
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadDeviceBlacklist:
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self):
|
||||
return ReadFileTool()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dev_random_blocked(self, tool):
|
||||
result = await tool.execute(path="/dev/random")
|
||||
assert "Error" in result
|
||||
assert "blocked" in result.lower() or "device" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dev_urandom_blocked(self, tool):
|
||||
result = await tool.execute(path="/dev/urandom")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dev_zero_blocked(self, tool):
|
||||
result = await tool.execute(path="/dev/zero")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proc_fd_blocked(self, tool):
|
||||
result = await tool.execute(path="/proc/self/fd/0")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_symlink_to_dev_zero_blocked(self, tmp_path):
|
||||
tool = ReadFileTool(workspace=tmp_path)
|
||||
link = tmp_path / "zero-link"
|
||||
link.symlink_to("/dev/zero")
|
||||
result = await tool.execute(path=str(link))
|
||||
assert "Error" in result
|
||||
assert "blocked" in result.lower() or "device" in result.lower()
|
||||
@@ -323,3 +323,27 @@ async def test_subagent_registers_grep_and_glob(tmp_path: Path) -> None:
|
||||
|
||||
assert "grep" in captured["tool_names"]
|
||||
assert "glob" in captured["tool_names"]
|
||||
|
||||
|
||||
def test_subagent_prompt_respects_disabled_skills(tmp_path: Path) -> None:
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
skills_dir = tmp_path / "skills"
|
||||
(skills_dir / "alpha").mkdir(parents=True)
|
||||
(skills_dir / "alpha" / "SKILL.md").write_text("# Alpha\n\nhidden\n", encoding="utf-8")
|
||||
(skills_dir / "beta").mkdir(parents=True)
|
||||
(skills_dir / "beta" / "SKILL.md").write_text("# Beta\n\nshown\n", encoding="utf-8")
|
||||
|
||||
mgr = SubagentManager(
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
bus=bus,
|
||||
max_tool_result_chars=4096,
|
||||
disabled_skills=["alpha"],
|
||||
)
|
||||
|
||||
prompt = mgr._build_subagent_prompt()
|
||||
|
||||
assert "alpha" not in prompt
|
||||
assert "beta" in prompt
|
||||
|
||||
@@ -120,6 +120,27 @@ async def test_jina_search(monkeypatch):
|
||||
assert "https://jina.ai" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kagi_search(monkeypatch):
|
||||
async def mock_get(self, url, **kw):
|
||||
assert "kagi.com/api/v0/search" in url
|
||||
assert kw["headers"]["Authorization"] == "Bot kagi-key"
|
||||
assert kw["params"] == {"q": "test", "limit": 2}
|
||||
return _response(json={
|
||||
"data": [
|
||||
{"t": 0, "title": "Kagi Result", "url": "https://kagi.com", "snippet": "Premium search"},
|
||||
{"t": 1, "list": ["ignored related search"]},
|
||||
]
|
||||
})
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)
|
||||
tool = _tool(provider="kagi", api_key="kagi-key")
|
||||
result = await tool.execute(query="test", count=2)
|
||||
assert "Kagi Result" in result
|
||||
assert "https://kagi.com" in result
|
||||
assert "ignored related search" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_provider():
|
||||
tool = _tool(provider="unknown")
|
||||
@@ -189,6 +210,23 @@ async def test_jina_422_falls_back_to_duckduckgo(monkeypatch):
|
||||
assert "DuckDuckGo fallback" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kagi_fallback_to_duckduckgo_when_no_key(monkeypatch):
|
||||
class MockDDGS:
|
||||
def __init__(self, **kw):
|
||||
pass
|
||||
|
||||
def text(self, query, max_results=5):
|
||||
return [{"title": "Fallback", "href": "https://ddg.example", "body": "DuckDuckGo fallback"}]
|
||||
|
||||
monkeypatch.setattr("ddgs.DDGS", MockDDGS)
|
||||
monkeypatch.delenv("KAGI_API_KEY", raising=False)
|
||||
|
||||
tool = _tool(provider="kagi", api_key="")
|
||||
result = await tool.execute(query="test")
|
||||
assert "Fallback" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jina_search_uses_path_encoded_query(monkeypatch):
|
||||
calls = {}
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
import pytest
|
||||
|
||||
from nanobot.utils.helpers import strip_think
|
||||
|
||||
|
||||
class TestStripThinkTag:
|
||||
"""Test <thought>...</thought> block stripping (Gemma 4 and similar models)."""
|
||||
|
||||
def test_closed_tag(self):
|
||||
assert strip_think("Hello <thought>reasoning</thought> World") == "Hello World"
|
||||
|
||||
def test_unclosed_trailing_tag(self):
|
||||
assert strip_think("<thought>ongoing...") == ""
|
||||
|
||||
def test_multiline_tag(self):
|
||||
assert strip_think("<thought>\nline1\nline2\n</thought>End") == "End"
|
||||
|
||||
def test_tag_with_nested_angle_brackets(self):
|
||||
text = "<thought>a < 3 and b > 2</thought>result"
|
||||
assert strip_think(text) == "result"
|
||||
|
||||
def test_multiple_tag_blocks(self):
|
||||
text = "A<thought>x</thought>B<thought>y</thought>C"
|
||||
assert strip_think(text) == "ABC"
|
||||
|
||||
def test_tag_only_whitespace_inside(self):
|
||||
assert strip_think("before<thought> </thought>after") == "beforeafter"
|
||||
|
||||
def test_self_closing_tag_not_matched(self):
|
||||
assert strip_think("<thought/>some text") == "<thought/>some text"
|
||||
|
||||
def test_normal_text_unchanged(self):
|
||||
assert strip_think("Just normal text") == "Just normal text"
|
||||
|
||||
def test_empty_string(self):
|
||||
assert strip_think("") == ""
|
||||
|
||||
|
||||
class TestStripThinkFalsePositive:
|
||||
"""Ensure mid-content <think>/<thought> tags are NOT stripped (#3004)."""
|
||||
|
||||
def test_backtick_think_tag_preserved(self):
|
||||
text = "*Think Stripping:* A new utility to strip `<think>` tags from output."
|
||||
assert strip_think(text) == text
|
||||
|
||||
def test_prose_think_tag_preserved(self):
|
||||
text = "The model emits <think> at the start of its response."
|
||||
assert strip_think(text) == text
|
||||
|
||||
def test_code_block_think_tag_preserved(self):
|
||||
text = "Example:\n```\ntext = re.sub(r\"<think>[\\s\\S]*\", \"\", text)\n```\nDone."
|
||||
assert strip_think(text) == text
|
||||
|
||||
def test_backtick_thought_tag_preserved(self):
|
||||
text = "Gemma 4 uses `<thought>` blocks for reasoning."
|
||||
assert strip_think(text) == text
|
||||
|
||||
def test_prefix_unclosed_think_still_stripped(self):
|
||||
assert strip_think("<think>reasoning without closing") == ""
|
||||
|
||||
def test_prefix_unclosed_think_with_whitespace(self):
|
||||
assert strip_think(" <think>reasoning...") == ""
|
||||
|
||||
def test_prefix_unclosed_thought_still_stripped(self):
|
||||
assert strip_think("<thought>reasoning without closing") == ""
|
||||
Reference in New Issue
Block a user