325 lines
11 KiB
Python
325 lines
11 KiB
Python
import asyncio
|
|
from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Literal, TypeAlias, cast
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from nanoid import generate as generate_nanoid
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src import crud, models, schemas
|
|
from src.utils.queue_payload import create_payload
|
|
from src.utils.work_unit import construct_work_unit_key
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
async def clean_queue_tables(db_session: AsyncSession) -> AsyncGenerator[None, None]:
|
|
"""Clean up queue-related tables before each test to ensure isolation.
|
|
|
|
This prevents webhook queue items and active queue sessions from previous tests
|
|
from polluting subsequent tests, which can cause issues when tests have a limit
|
|
on how many work units can be claimed (e.g., WORKERS=1).
|
|
"""
|
|
# Clean up before the test
|
|
await db_session.execute(delete(models.ActiveQueueSession))
|
|
await db_session.execute(delete(models.QueueItem))
|
|
await db_session.commit()
|
|
|
|
yield
|
|
|
|
|
|
QueuePayload: TypeAlias = dict[str, Any]
|
|
QueuePayloadEntry: TypeAlias = QueuePayload | tuple[QueuePayload, int | None]
|
|
|
|
|
|
@pytest.fixture
|
|
async def sample_session_with_peers(
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> tuple[models.Session, list[models.Peer]]:
|
|
"""Create a sample session with multiple peers for testing deriver functionality"""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create additional peers
|
|
peer2 = models.Peer(name=str(generate_nanoid()), workspace_name=workspace.name)
|
|
peer3 = models.Peer(name=str(generate_nanoid()), workspace_name=workspace.name)
|
|
db_session.add_all([peer2, peer3])
|
|
await db_session.flush()
|
|
|
|
# Create session with peer configurations
|
|
session = (
|
|
await crud.get_or_create_session(
|
|
db_session,
|
|
schemas.SessionCreate(
|
|
name=str(generate_nanoid()),
|
|
peers={
|
|
peer1.name: schemas.SessionPeerConfig(observe_me=True),
|
|
peer2.name: schemas.SessionPeerConfig(observe_others=True),
|
|
peer3.name: schemas.SessionPeerConfig(), # No special observation settings
|
|
},
|
|
),
|
|
workspace.name,
|
|
)
|
|
).resource
|
|
await db_session.commit()
|
|
|
|
return session, [peer1, peer2, peer3]
|
|
|
|
|
|
@pytest.fixture
|
|
async def sample_messages(
|
|
db_session: AsyncSession,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
) -> list[models.Message]:
|
|
"""Create sample messages for testing deriver functionality"""
|
|
session, peers = sample_session_with_peers
|
|
peer1, peer2, peer3 = peers
|
|
|
|
# Create multiple messages from different peers
|
|
messages_data = [
|
|
{
|
|
"session_name": session.name,
|
|
"content": "Hello, this is the first message from peer1",
|
|
"peer_name": peer1.name,
|
|
"workspace_name": session.workspace_name,
|
|
"seq_in_session": 1,
|
|
},
|
|
{
|
|
"session_name": session.name,
|
|
"content": "Hi there! This is a response from peer2",
|
|
"peer_name": peer2.name,
|
|
"workspace_name": session.workspace_name,
|
|
"seq_in_session": 2,
|
|
},
|
|
{
|
|
"session_name": session.name,
|
|
"content": "I'm just observing this conversation as peer3",
|
|
"peer_name": peer3.name,
|
|
"workspace_name": session.workspace_name,
|
|
"seq_in_session": 3,
|
|
},
|
|
]
|
|
|
|
messages: list[models.Message] = []
|
|
for msg_data in messages_data:
|
|
message = models.Message(**msg_data)
|
|
db_session.add(message)
|
|
messages.append(message)
|
|
|
|
await db_session.commit()
|
|
|
|
# Query the messages again to get the committed versions
|
|
result = await db_session.execute(
|
|
select(models.Message)
|
|
.where(models.Message.session_name == session.name)
|
|
.order_by(models.Message.id)
|
|
)
|
|
messages = list(result.scalars().all())
|
|
|
|
return messages
|
|
|
|
|
|
@pytest.fixture
|
|
def create_queue_payload() -> Callable[..., Any]:
|
|
"""Helper function to create queue payloads for testing"""
|
|
|
|
def _create_payload(
|
|
message: models.Message,
|
|
task_type: Literal["representation", "summary"],
|
|
observer: str | None = None,
|
|
observed: str | None = None,
|
|
message_seq_in_session: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Create a queue payload for testing"""
|
|
message_dict = {
|
|
"workspace_name": message.workspace_name,
|
|
"session_name": message.session_name,
|
|
"message_id": message.id,
|
|
"content": message.content,
|
|
"created_at": message.created_at or datetime.now(timezone.utc),
|
|
"message_public_id": message.public_id,
|
|
}
|
|
|
|
configuration = schemas.ResolvedConfiguration(
|
|
reasoning=schemas.ResolvedReasoningConfiguration(enabled=True),
|
|
peer_card=schemas.ResolvedPeerCardConfiguration(use=True, create=True),
|
|
summary=schemas.ResolvedSummaryConfiguration(
|
|
enabled=True,
|
|
messages_per_short_summary=10,
|
|
messages_per_long_summary=20,
|
|
),
|
|
dream=schemas.ResolvedDreamConfiguration(enabled=True),
|
|
)
|
|
|
|
return create_payload(
|
|
message=message_dict,
|
|
configuration=configuration,
|
|
task_type=task_type,
|
|
message_seq_in_session=message_seq_in_session,
|
|
observers=[observer] if observer else None,
|
|
observed=observed,
|
|
)
|
|
|
|
return _create_payload
|
|
|
|
|
|
@pytest.fixture
|
|
async def add_queue_items(
|
|
db_session: AsyncSession,
|
|
) -> Callable[
|
|
[Sequence[QueuePayloadEntry], str, str], Awaitable[list[models.QueueItem]]
|
|
]:
|
|
"""Helper function to add queue items to the database"""
|
|
|
|
async def _add_items(
|
|
payloads: Sequence[QueuePayloadEntry],
|
|
session_id: str,
|
|
workspace_name: str,
|
|
) -> list[models.QueueItem]:
|
|
"""Add queue items to the database and return them"""
|
|
queue_items: list[models.QueueItem] = []
|
|
for payload_entry in payloads:
|
|
payload: QueuePayload
|
|
message_id: int | None
|
|
if isinstance(payload_entry, tuple):
|
|
payload, message_id = payload_entry
|
|
else:
|
|
payload = payload_entry
|
|
message_id = cast(int | None, payload.get("message_id"))
|
|
# Generate work_unit_key from the payload
|
|
task_type = cast(str, payload.get("task_type", "unknown"))
|
|
work_unit_key = construct_work_unit_key(workspace_name, payload)
|
|
|
|
queue_item = models.QueueItem(
|
|
session_id=session_id,
|
|
task_type=task_type,
|
|
work_unit_key=work_unit_key,
|
|
payload=payload,
|
|
processed=False,
|
|
workspace_name=workspace_name,
|
|
message_id=message_id,
|
|
)
|
|
db_session.add(queue_item)
|
|
queue_items.append(queue_item)
|
|
|
|
await db_session.commit()
|
|
|
|
# Refresh to get the actual IDs
|
|
for item in queue_items:
|
|
await db_session.refresh(item)
|
|
|
|
return queue_items
|
|
|
|
return _add_items
|
|
|
|
|
|
@pytest.fixture
|
|
async def sample_queue_items(
|
|
db_session: AsyncSession, # noqa: ARG001 # pyright: ignore[reportUnusedParameter]
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
sample_messages: list[models.Message],
|
|
create_queue_payload: Callable[..., Any],
|
|
add_queue_items: Callable[..., Any],
|
|
) -> list[models.QueueItem]:
|
|
"""Create sample queue items for testing"""
|
|
session, peers = sample_session_with_peers
|
|
_peer1, peer2, _peer3 = peers
|
|
messages = sample_messages
|
|
|
|
# Create various types of queue payloads
|
|
payloads: list[tuple[dict[str, Any], int]] = []
|
|
|
|
# Create representation payloads for each message
|
|
for message in messages:
|
|
# Self-representation (peer observing themselves)
|
|
payload1 = create_queue_payload(
|
|
message=message,
|
|
task_type="representation",
|
|
observer=message.peer_name,
|
|
observed=message.peer_name,
|
|
)
|
|
payloads.append((payload1, message.id))
|
|
|
|
# Representation for observer peer
|
|
payload2 = create_queue_payload(
|
|
message=message,
|
|
task_type="representation",
|
|
observer=peer2.name, # peer2 observes others
|
|
observed=message.peer_name,
|
|
)
|
|
payloads.append((payload2, message.id))
|
|
|
|
# Create summary payloads for session
|
|
for i, message in enumerate(messages):
|
|
payload = create_queue_payload(
|
|
message=message,
|
|
task_type="summary",
|
|
message_seq_in_session=i + 1,
|
|
)
|
|
payloads.append((payload, message.id))
|
|
|
|
# Add all payloads as queue items
|
|
queue_items = await add_queue_items(payloads, session.id, session.workspace_name)
|
|
|
|
return queue_items
|
|
|
|
|
|
@pytest.fixture
|
|
async def create_active_queue_session(db_session: AsyncSession) -> Callable[..., Any]:
|
|
"""Helper function to create active queue sessions for testing work unit tracking"""
|
|
|
|
async def _create_active_session(
|
|
work_unit_key: str,
|
|
) -> models.ActiveQueueSession:
|
|
"""Create an active queue session"""
|
|
active_session = models.ActiveQueueSession(
|
|
work_unit_key=work_unit_key,
|
|
)
|
|
db_session.add(active_session)
|
|
await db_session.commit()
|
|
await db_session.refresh(active_session)
|
|
return active_session
|
|
|
|
return _create_active_session
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_queue_manager(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: # pyright: ignore[reportUnusedParameter]
|
|
"""Mock the queue manager to avoid actual queue processing"""
|
|
from src.deriver.queue_manager import QueueManager
|
|
|
|
# Create a mock queue manager
|
|
mock_manager = AsyncMock(spec=QueueManager)
|
|
|
|
# Mock the methods we might need to test
|
|
mock_manager.initialize = AsyncMock()
|
|
mock_manager.shutdown = AsyncMock()
|
|
mock_manager.process_work_unit = AsyncMock()
|
|
mock_manager.get_available_work_units = AsyncMock(return_value=[])
|
|
mock_manager.add_task = MagicMock()
|
|
mock_manager.track_work_unit = MagicMock()
|
|
mock_manager.untrack_work_unit = MagicMock()
|
|
|
|
# Mock the attributes
|
|
mock_manager.shutdown_event = asyncio.Event()
|
|
mock_manager.active_tasks = set()
|
|
mock_manager.owned_work_units = set()
|
|
mock_manager.queue_empty_flag = asyncio.Event()
|
|
mock_manager.workers = 1
|
|
mock_manager.semaphore = asyncio.Semaphore(1)
|
|
|
|
return mock_manager
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_representation_manager(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: # pyright: ignore[reportUnusedParameter]
|
|
"""Mock the representation manager to avoid actual embedding operations"""
|
|
from src.crud.representation import RepresentationManager
|
|
|
|
mock_manager = AsyncMock(spec=RepresentationManager)
|
|
mock_manager.save_representation.return_value = 0
|
|
|
|
return mock_manager
|