660 lines
24 KiB
Python
660 lines
24 KiB
Python
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src import models
|
|
from src.config import settings
|
|
from src.deriver.queue_manager import QueueManager
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestQueueProcessing:
|
|
"""Test suite for queue processing functionality"""
|
|
|
|
async def test_get_and_claim_work_units(
|
|
self,
|
|
sample_queue_items: list[models.QueueItem],
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
) -> None:
|
|
"""Test that get_and_claim_work_units correctly identifies unprocessed work"""
|
|
session, _peers = sample_session_with_peers # pyright: ignore[reportUnusedVariable]
|
|
|
|
# Verify we have queue items from our test setup
|
|
assert len(sample_queue_items) == 9 # 6 representation + 3 summary
|
|
|
|
# Create a queue manager instance
|
|
queue_manager = QueueManager()
|
|
|
|
# Get available work units
|
|
work_units = await queue_manager.get_and_claim_work_units()
|
|
|
|
# Should have some work units available (may include items from other tests)
|
|
assert len(work_units) > 0
|
|
|
|
# Check that all work units have the expected structure
|
|
for work_unit in work_units:
|
|
assert isinstance(work_unit, str)
|
|
assert work_unit.split(":")[0] in ["representation", "summary"]
|
|
|
|
# The test is mainly verifying that get_and_claim_work_units works without errors
|
|
# and returns properly structured work unit key strings
|
|
|
|
async def test_work_unit_claiming(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter]
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
) -> None:
|
|
"""Test that work units can be claimed and are not available to other workers"""
|
|
_session, _peers = sample_session_with_peers
|
|
|
|
# Create a queue manager instance
|
|
queue_manager = QueueManager()
|
|
|
|
# Get available work units
|
|
work_units = await queue_manager.get_and_claim_work_units()
|
|
assert len(work_units) > 0
|
|
|
|
# The API already claimed returned units; verify it's tracked and not returned again
|
|
from sqlalchemy import select
|
|
|
|
work_unit = work_units[0]
|
|
tracked = (
|
|
await db_session.execute(
|
|
select(models.ActiveQueueSession).where(
|
|
models.ActiveQueueSession.work_unit_key == work_unit
|
|
)
|
|
)
|
|
).scalar_one_or_none()
|
|
assert tracked is not None
|
|
|
|
# Get available work units again - the claimed one should not be available
|
|
remaining_work_units = await queue_manager.get_and_claim_work_units()
|
|
|
|
# The claimed work unit should not be in the remaining list
|
|
assert work_unit not in remaining_work_units
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_and_claim_excludes_already_claimed(
|
|
self,
|
|
sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter]
|
|
) -> None:
|
|
queue_manager = QueueManager()
|
|
first_batch = await queue_manager.get_and_claim_work_units()
|
|
assert len(first_batch) > 0
|
|
|
|
# Call again; previously claimed keys should not appear
|
|
second_batch = await queue_manager.get_and_claim_work_units()
|
|
assert all(k not in second_batch for k in first_batch)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_claim_work_unit_conflict_returns_false(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter]
|
|
) -> None:
|
|
# Pre-create an active session for a key
|
|
queue_manager = QueueManager()
|
|
claimed = await queue_manager.get_and_claim_work_units()
|
|
assert len(claimed) > 0
|
|
key = claimed[0]
|
|
|
|
# Trying to claim the same key again via the API should return empty list
|
|
claimed_again = await queue_manager.claim_work_units(db_session, [key])
|
|
assert claimed_again == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_next_message_orders_and_filters_simple(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
create_queue_payload: Callable[..., Any],
|
|
add_queue_items: Callable[..., Any],
|
|
) -> None:
|
|
from sqlalchemy import select
|
|
|
|
session, peers = sample_session_with_peers
|
|
peer = peers[0]
|
|
|
|
# Create and save messages to the database first
|
|
messages: list[models.Message] = []
|
|
for _ in range(3):
|
|
message = models.Message(
|
|
session_name=session.name,
|
|
workspace_name=session.workspace_name,
|
|
peer_name=peer.name,
|
|
content="hello",
|
|
token_count=10,
|
|
)
|
|
db_session.add(message)
|
|
messages.append(message)
|
|
|
|
await db_session.commit()
|
|
|
|
# Refresh to get the actual IDs
|
|
for message in messages:
|
|
await db_session.refresh(message)
|
|
|
|
payloads: list[Any] = []
|
|
for message in messages:
|
|
payload = create_queue_payload( # type: ignore[reportUnknownArgumentType]
|
|
message=message,
|
|
task_type="representation",
|
|
sender_name=peer.name,
|
|
target_name=peer.name,
|
|
)
|
|
payloads.append(payload)
|
|
|
|
items = await add_queue_items(payloads, session.id)
|
|
# Determine ascending order by DB id
|
|
ordered = (
|
|
(
|
|
await db_session.execute(
|
|
select(models.QueueItem)
|
|
.where(models.QueueItem.work_unit_key == items[0].work_unit_key)
|
|
.order_by(models.QueueItem.id)
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
first, second = ordered[0], ordered[1]
|
|
|
|
qm = QueueManager()
|
|
batch = await qm.get_message_batch(
|
|
first.work_unit_key,
|
|
task_type="representation",
|
|
)
|
|
nxt = batch[0] if batch else None
|
|
assert nxt is not None and nxt.id == first.id
|
|
|
|
# Mark first processed, next should be the second
|
|
first.processed = True
|
|
await db_session.commit()
|
|
batch2 = await qm.get_message_batch(
|
|
first.work_unit_key,
|
|
task_type="representation",
|
|
)
|
|
nxt2 = batch2[0] if batch2 else None
|
|
assert nxt2 is not None and nxt2.id == second.id
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_work_unit_removes_row(
|
|
self,
|
|
sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter]
|
|
db_session: AsyncSession,
|
|
) -> None:
|
|
from sqlalchemy import select
|
|
|
|
qm = QueueManager()
|
|
claimed = await qm.get_and_claim_work_units()
|
|
assert len(claimed) > 0
|
|
key = claimed[0]
|
|
|
|
removed = await qm._cleanup_work_unit(key) # pyright: ignore[reportPrivateUsage]
|
|
assert removed is True
|
|
|
|
remaining = (
|
|
await db_session.execute(
|
|
select(models.ActiveQueueSession).where(
|
|
models.ActiveQueueSession.work_unit_key == key
|
|
)
|
|
)
|
|
).scalar_one_or_none()
|
|
assert remaining is None
|
|
|
|
async def test_stale_work_unit_cleanup(
|
|
self,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
) -> None:
|
|
"""Test that stale work units are cleaned up properly"""
|
|
_session, _peers = sample_session_with_peers
|
|
|
|
# Create an active queue session with an old timestamp
|
|
# from datetime import datetime, timedelta, timezone
|
|
|
|
# datetime.now(timezone.utc) - timedelta(minutes=10)
|
|
|
|
# We'll test this by checking that the cleanup logic works in get_and_claim_work_units
|
|
# which is called by the queue manager during normal operation
|
|
queue_manager = QueueManager()
|
|
|
|
# Get available work units - this should clean up stale entries
|
|
work_units = await queue_manager.get_and_claim_work_units()
|
|
|
|
# This test ensures the cleanup logic doesn't break, though we don't have stale entries yet
|
|
assert isinstance(work_units, list)
|
|
|
|
async def test_work_unit_key_format(
|
|
self, sample_session_with_peers: tuple[models.Session, list[models.Peer]]
|
|
) -> None:
|
|
"""Test that work unit keys have the correct format"""
|
|
session, peers = sample_session_with_peers
|
|
peer1, peer2, _ = peers
|
|
|
|
# Create a representation work unit key
|
|
# Format: task_type:workspace:session:sender:target
|
|
work_unit_key = (
|
|
f"representation:workspace1:{session.name}:{peer1.name}:{peer2.name}"
|
|
)
|
|
|
|
# Check that the key contains the expected information
|
|
assert session.name in work_unit_key
|
|
assert peer1.name in work_unit_key
|
|
assert peer2.name in work_unit_key
|
|
assert "representation" in work_unit_key
|
|
assert "workspace1" in work_unit_key
|
|
|
|
# Create a summary work unit key
|
|
# Summary work units use None for sender/target
|
|
summary_work_unit_key = f"summary:workspace1:{session.name}:None:None"
|
|
|
|
assert session.name in summary_work_unit_key
|
|
assert "None" in summary_work_unit_key
|
|
assert "summary" in summary_work_unit_key
|
|
assert "workspace1" in summary_work_unit_key
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_representation_batching_respects_token_limits(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
create_queue_payload: Callable[..., Any],
|
|
) -> None:
|
|
"""Test that representation tasks are batched based on token limits"""
|
|
from unittest.mock import patch
|
|
|
|
session, peers = sample_session_with_peers
|
|
peer = peers[0]
|
|
|
|
# Create messages with token counts that exceed batch limit
|
|
limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
|
|
token_counts = [limit // 2, limit // 2, limit // 2]
|
|
|
|
# Create and save messages to the database first
|
|
messages: list[models.Message] = []
|
|
for i, token_count in enumerate(token_counts):
|
|
message = models.Message(
|
|
session_name=session.name,
|
|
workspace_name=session.workspace_name,
|
|
peer_name=peer.name,
|
|
content=f"Test message {i}",
|
|
token_count=token_count,
|
|
)
|
|
db_session.add(message)
|
|
messages.append(message)
|
|
|
|
await db_session.commit()
|
|
|
|
# Refresh to get the actual IDs
|
|
for message in messages:
|
|
await db_session.refresh(message)
|
|
|
|
# Create queue items with token counts
|
|
payloads = [
|
|
create_queue_payload( # type: ignore[reportUnknownArgumentType]
|
|
message=msg,
|
|
task_type="representation",
|
|
sender_name=peer.name,
|
|
target_name=peer.name,
|
|
)
|
|
for msg in messages
|
|
]
|
|
|
|
# Add items with token counts
|
|
from src.deriver.utils import get_work_unit_key
|
|
|
|
queue_items: list[models.QueueItem] = []
|
|
for payload in payloads:
|
|
task_type = payload.get("task_type", "unknown")
|
|
work_unit_key = get_work_unit_key(task_type, payload)
|
|
|
|
queue_item = models.QueueItem(
|
|
session_id=session.id,
|
|
task_type=task_type,
|
|
work_unit_key=work_unit_key,
|
|
payload=payload,
|
|
processed=False,
|
|
)
|
|
db_session.add(queue_item)
|
|
queue_items.append(queue_item)
|
|
|
|
await db_session.commit()
|
|
for item in queue_items:
|
|
await db_session.refresh(item)
|
|
|
|
# Mock process_items to capture batches
|
|
processed_batches: list[dict[str, Any]] = []
|
|
|
|
async def mock_process_items(
|
|
task_type: str, queue_payloads: list[dict[str, Any]]
|
|
) -> None:
|
|
processed_batches.append(
|
|
{
|
|
"task_type": task_type,
|
|
"payload_count": len(queue_payloads),
|
|
}
|
|
)
|
|
|
|
# Process work unit and verify batching
|
|
qm = QueueManager()
|
|
with patch(
|
|
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
|
|
):
|
|
await qm.process_work_unit(queue_items[0].work_unit_key)
|
|
|
|
# Should create 2 batches due to token limits
|
|
assert len(processed_batches) == 2
|
|
assert processed_batches[0]["payload_count"] == 2
|
|
assert processed_batches[1]["payload_count"] == 1
|
|
assert all(b["task_type"] == "representation" for b in processed_batches)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_message_processing(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
create_queue_payload: Callable[..., Any],
|
|
) -> None:
|
|
"""Test that multiple summary messages in same work unit are processed separately"""
|
|
from unittest.mock import patch
|
|
|
|
session, peers = sample_session_with_peers
|
|
peer = peers[0]
|
|
|
|
# Create two summary messages
|
|
token_counts = [500, 600]
|
|
messages = [
|
|
models.Message(
|
|
id=999,
|
|
session_name=session.name,
|
|
workspace_name=session.workspace_name,
|
|
peer_name=peer.name,
|
|
content="First summary message",
|
|
),
|
|
models.Message(
|
|
id=1000,
|
|
session_name=session.name,
|
|
workspace_name=session.workspace_name,
|
|
peer_name=peer.name,
|
|
content="Second summary message",
|
|
),
|
|
]
|
|
|
|
# Create payloads and queue items
|
|
queue_items: list[models.QueueItem] = []
|
|
for i, message in enumerate(messages):
|
|
payload = create_queue_payload(
|
|
message, "summary", message_seq_in_session=i + 1
|
|
)
|
|
payload["token_count"] = token_counts[i]
|
|
from src.deriver.utils import get_work_unit_key
|
|
|
|
work_unit_key = get_work_unit_key("summary", payload)
|
|
|
|
queue_item = models.QueueItem(
|
|
session_id=session.id,
|
|
task_type="summary",
|
|
work_unit_key=work_unit_key,
|
|
payload=payload,
|
|
processed=False,
|
|
)
|
|
db_session.add(queue_item)
|
|
queue_items.append(queue_item)
|
|
|
|
await db_session.commit()
|
|
|
|
# Mock and process work unit
|
|
processed_batches: list[dict[str, Any]] = []
|
|
|
|
async def mock_process_items(
|
|
task_type: str, queue_payloads: list[dict[str, Any]]
|
|
) -> None:
|
|
processed_batches.append(
|
|
{"task_type": task_type, "payload_count": len(queue_payloads)}
|
|
)
|
|
|
|
qm = QueueManager()
|
|
work_unit_key = queue_items[0].work_unit_key
|
|
with patch(
|
|
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
|
|
):
|
|
await qm.process_work_unit(work_unit_key)
|
|
|
|
# Verify both messages were processed in separate batches
|
|
assert len(processed_batches) == 2
|
|
assert all(batch["task_type"] == "summary" for batch in processed_batches)
|
|
assert all(batch["payload_count"] == 1 for batch in processed_batches)
|
|
|
|
# Verify the corresponding DB records are marked as processed
|
|
from sqlalchemy import select
|
|
|
|
# Query for the summary queue items that were processed
|
|
processed_items = (
|
|
(
|
|
await db_session.execute(
|
|
select(models.QueueItem)
|
|
.where(models.QueueItem.work_unit_key == work_unit_key)
|
|
.where(models.QueueItem.task_type == "summary")
|
|
.order_by(models.QueueItem.id)
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
|
|
# Assert we found both summary items
|
|
assert len(processed_items) == 2
|
|
|
|
# Assert both items are marked as processed
|
|
assert all(item.processed is True for item in processed_items)
|
|
|
|
# Optionally verify the items have the expected token counts from the messages
|
|
expected_token_counts = [500, 600] # From the test messages
|
|
actual_token_counts = [
|
|
item.payload.get("token_count") or 0 for item in processed_items
|
|
]
|
|
assert sorted(actual_token_counts) == sorted(expected_token_counts)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_first_message_exceeds_token_limit_still_included(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
create_queue_payload: Callable[..., Any],
|
|
) -> None:
|
|
"""Test that if the first message exceeds BATCH_MAX_TOKENS, it's still included alone"""
|
|
from unittest.mock import patch
|
|
|
|
session, peers = sample_session_with_peers
|
|
peer = peers[0]
|
|
|
|
# Create messages where first message exceeds the batch limit
|
|
limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
|
|
token_counts = [limit + 1000, 100, 200] # First message way over limit
|
|
|
|
# Create and save messages to the database first
|
|
messages: list[models.Message] = []
|
|
for i, token_count in enumerate(token_counts):
|
|
message = models.Message(
|
|
session_name=session.name,
|
|
workspace_name=session.workspace_name,
|
|
peer_name=peer.name,
|
|
content=f"Test message {i}",
|
|
token_count=token_count,
|
|
)
|
|
db_session.add(message)
|
|
messages.append(message)
|
|
|
|
await db_session.commit()
|
|
|
|
# Refresh to get the actual IDs
|
|
for message in messages:
|
|
await db_session.refresh(message)
|
|
|
|
# Create queue items
|
|
payloads = [
|
|
create_queue_payload( # type: ignore[reportUnknownArgumentType]
|
|
message=msg,
|
|
task_type="representation",
|
|
sender_name=peer.name,
|
|
target_name=peer.name,
|
|
)
|
|
for msg in messages
|
|
]
|
|
|
|
# Add items to queue
|
|
from src.deriver.utils import get_work_unit_key
|
|
|
|
queue_items: list[models.QueueItem] = []
|
|
for payload in payloads:
|
|
task_type = payload.get("task_type", "unknown")
|
|
work_unit_key = get_work_unit_key(task_type, payload)
|
|
|
|
queue_item = models.QueueItem(
|
|
session_id=session.id,
|
|
task_type=task_type,
|
|
work_unit_key=work_unit_key,
|
|
payload=payload,
|
|
processed=False,
|
|
)
|
|
db_session.add(queue_item)
|
|
queue_items.append(queue_item)
|
|
|
|
await db_session.commit()
|
|
for item in queue_items:
|
|
await db_session.refresh(item)
|
|
|
|
# Mock process_items to capture batches
|
|
processed_batches: list[dict[str, Any]] = []
|
|
|
|
async def mock_process_items(
|
|
task_type: str, queue_payloads: list[dict[str, Any]]
|
|
) -> None:
|
|
processed_batches.append(
|
|
{
|
|
"task_type": task_type,
|
|
"payload_count": len(queue_payloads),
|
|
}
|
|
)
|
|
|
|
# Process work unit and verify batching
|
|
qm = QueueManager()
|
|
with patch(
|
|
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
|
|
):
|
|
await qm.process_work_unit(queue_items[0].work_unit_key)
|
|
|
|
# Should create 2 batches: first large message alone, then second and third together
|
|
assert len(processed_batches) == 2
|
|
assert (
|
|
processed_batches[0]["payload_count"] == 1
|
|
) # First message (over limit) alone
|
|
assert processed_batches[1]["payload_count"] == 2 # Second and third messages
|
|
assert all(b["task_type"] == "representation" for b in processed_batches)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_message_exactly_at_token_limit(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
|
create_queue_payload: Callable[..., Any],
|
|
) -> None:
|
|
"""Test boundary condition when cumulative sum exactly equals limit"""
|
|
from unittest.mock import patch
|
|
|
|
session, peers = sample_session_with_peers
|
|
peer = peers[0]
|
|
|
|
# Create messages that test the exact boundary
|
|
limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
|
|
token_counts = [
|
|
limit // 2,
|
|
limit // 2,
|
|
1,
|
|
] # First two exactly at limit, third exceeds
|
|
|
|
# Create and save messages to the database first
|
|
messages: list[models.Message] = []
|
|
for i, token_count in enumerate(token_counts):
|
|
message = models.Message(
|
|
session_name=session.name,
|
|
workspace_name=session.workspace_name,
|
|
peer_name=peer.name,
|
|
content=f"Test message {i}",
|
|
token_count=token_count,
|
|
)
|
|
db_session.add(message)
|
|
messages.append(message)
|
|
|
|
await db_session.commit()
|
|
|
|
# Refresh to get the actual IDs
|
|
for message in messages:
|
|
await db_session.refresh(message)
|
|
|
|
# Create queue items
|
|
payloads = [
|
|
create_queue_payload( # type: ignore[reportUnknownArgumentType]
|
|
message=msg,
|
|
task_type="representation",
|
|
sender_name=peer.name,
|
|
target_name=peer.name,
|
|
)
|
|
for msg in messages
|
|
]
|
|
|
|
# Add items to queue
|
|
from src.deriver.utils import get_work_unit_key
|
|
|
|
queue_items: list[models.QueueItem] = []
|
|
for payload in payloads:
|
|
task_type = payload.get("task_type", "unknown")
|
|
work_unit_key = get_work_unit_key(task_type, payload)
|
|
|
|
queue_item = models.QueueItem(
|
|
session_id=session.id,
|
|
task_type=task_type,
|
|
work_unit_key=work_unit_key,
|
|
payload=payload,
|
|
processed=False,
|
|
)
|
|
db_session.add(queue_item)
|
|
queue_items.append(queue_item)
|
|
|
|
await db_session.commit()
|
|
for item in queue_items:
|
|
await db_session.refresh(item)
|
|
|
|
# Mock process_items to capture batches
|
|
processed_batches: list[dict[str, Any]] = []
|
|
|
|
async def mock_process_items(
|
|
task_type: str, queue_payloads: list[dict[str, Any]]
|
|
) -> None:
|
|
processed_batches.append(
|
|
{
|
|
"task_type": task_type,
|
|
"payload_count": len(queue_payloads),
|
|
}
|
|
)
|
|
|
|
# Process work unit and verify batching
|
|
qm = QueueManager()
|
|
with patch(
|
|
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
|
|
):
|
|
await qm.process_work_unit(queue_items[0].work_unit_key)
|
|
|
|
# Should create 2 batches: first two messages together (exactly at limit), third alone
|
|
assert len(processed_batches) == 2
|
|
assert (
|
|
processed_batches[0]["payload_count"] == 2
|
|
) # First two messages (exactly at limit)
|
|
assert (
|
|
processed_batches[1]["payload_count"] == 1
|
|
) # Third message (exceeds limit)
|
|
assert all(b["task_type"] == "representation" for b in processed_batches)
|