"""Testing utilities for making assertions about scheduled tasks.
Example usage:
from docket import Docket, testing
docket = Docket("redis://localhost:6379/0")
# Schedule a task
await docket.add(my_task)("arg1", kwarg1="value1")
# Assert it's scheduled
await testing.assert_task_scheduled(docket, my_task, args=("arg1",))
# After completion
await worker.run_until_finished()
await testing.assert_no_tasks(docket)
"""
from collections.abc import Callable
from typing import Any
from docket.docket import Docket
from docket.execution import Execution
def _matches_criteria(
execution: Execution,
function: str | Callable[..., Any],
args: tuple[Any, ...] | None,
kwargs: dict[str, Any] | None,
key: str | None,
) -> bool:
"""Check if an execution matches the given criteria."""
# Check function name
function_name = function if isinstance(function, str) else function.__name__
if execution.function_name != function_name:
return False
# Check key if specified
if key is not None and execution.key != key:
return False
# Check args if specified
if args is not None and execution.args != args:
return False
# Check kwargs if specified (subset matching)
if kwargs is not None:
for k, v in kwargs.items():
if k not in execution.kwargs or execution.kwargs[k] != v:
return False
return True
def _format_criteria(
function: str | Callable[..., Any],
args: tuple[Any, ...] | None,
kwargs: dict[str, Any] | None,
key: str | None,
) -> str:
"""Format criteria for error messages."""
parts: list[str] = []
function_name = function if isinstance(function, str) else function.__name__
parts.append(f"function={function_name}")
if key is not None:
parts.append(f"key={key!r}")
if args is not None:
parts.append(f"args={args!r}")
if kwargs is not None:
parts.append(f"kwargs={kwargs!r}")
return ", ".join(parts)
async def assert_task_scheduled(
docket: Docket,
function: str | Callable[..., Any],
*,
args: tuple[Any, ...] | None = None,
kwargs: dict[str, Any] | None = None,
key: str | None = None,
) -> None:
"""Assert that a task matching the criteria is scheduled.
Args:
docket: The Docket instance to check
function: The task function or function name (string)
args: Optional tuple of positional arguments to match
kwargs: Optional dict of keyword arguments to match (subset matching)
key: Optional task key to match
Raises:
AssertionError: If no matching task is found
Example:
await assert_task_scheduled(docket, my_task)
await assert_task_scheduled(docket, my_task, args=("foo",))
await assert_task_scheduled(docket, "my_task", key="task-123")
"""
snapshot = await docket.snapshot()
# Check all scheduled tasks (both immediate and future)
all_tasks = list(snapshot.future)
for execution in all_tasks:
if _matches_criteria(execution, function, args, kwargs, key):
return
# Build error message
criteria = _format_criteria(function, args, kwargs, key)
function_name = function if isinstance(function, str) else function.__name__
if not all_tasks:
raise AssertionError(
f"Task {function_name} not found: no tasks scheduled on docket"
)
# Show what we found instead
found_tasks = [
f" - {e.function_name}(args={e.args!r}, kwargs={e.kwargs!r}, key={e.key!r})"
for e in all_tasks
]
found_str = "\n".join(found_tasks)
raise AssertionError(
f"Task {function_name} not found with {criteria}\n\n"
f"Scheduled tasks:\n{found_str}"
)
async def assert_task_not_scheduled(
docket: Docket,
function: str | Callable[..., Any],
*,
args: tuple[Any, ...] | None = None,
kwargs: dict[str, Any] | None = None,
key: str | None = None,
) -> None:
"""Assert that no task matching the criteria is scheduled.
Args:
docket: The Docket instance to check
function: The task function or function name (string)
args: Optional tuple of positional arguments to match
kwargs: Optional dict of keyword arguments to match (subset matching)
key: Optional task key to match
Raises:
AssertionError: If a matching task is found
Example:
await assert_task_not_scheduled(docket, my_task)
await assert_task_not_scheduled(docket, my_task, args=("foo",))
"""
snapshot = await docket.snapshot()
# Check all scheduled tasks (both immediate and future)
all_tasks = list(snapshot.future)
for execution in all_tasks:
if _matches_criteria(execution, function, args, kwargs, key):
function_name = function if isinstance(function, str) else function.__name__
raise AssertionError(
f"Task {function_name} found but should not be scheduled\n"
f"Found: {execution.function_name}(args={execution.args!r}, "
f"kwargs={execution.kwargs!r}, key={execution.key!r})"
)
async def assert_task_count(
docket: Docket,
function: str | Callable[..., Any] | None = None,
*,
count: int,
) -> None:
"""Assert the number of scheduled tasks matches the expected count.
Args:
docket: The Docket instance to check
function: Optional task function or name to count (if None, counts all tasks)
count: Expected number of tasks
Raises:
AssertionError: If the count doesn't match
Example:
await assert_task_count(docket, count=5) # All tasks
await assert_task_count(docket, my_task, count=2) # Specific function
"""
snapshot = await docket.snapshot()
# Check all scheduled tasks (both immediate and future)
all_tasks = list(snapshot.future)
if function is None:
actual_count = len(all_tasks)
function_desc = "all tasks"
else:
function_name = function if isinstance(function, str) else function.__name__
actual_count = sum(1 for e in all_tasks if e.function_name == function_name)
function_desc = f"tasks for {function_name}"
if actual_count != count:
raise AssertionError(f"Expected {count} {function_desc}, found {actual_count}")
async def assert_no_tasks(docket: Docket) -> None:
"""Assert that no tasks are scheduled on the docket.
Args:
docket: The Docket instance to check
Raises:
AssertionError: If any tasks are scheduled
Example:
await assert_no_tasks(docket)
"""
snapshot = await docket.snapshot()
# Check all scheduled tasks (both immediate and future)
all_tasks = list(snapshot.future)
if all_tasks:
found_tasks = [
f" - {e.function_name}(args={e.args!r}, kwargs={e.kwargs!r}, key={e.key!r})"
for e in all_tasks
]
found_str = "\n".join(found_tasks)
raise AssertionError(
f"Expected no tasks, found {len(all_tasks)} task(s) scheduled:\n{found_str}"
)
