Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions .github/workflows/codspeed.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
name: CodSpeed

on:
push:
branches: [ main ]
paths-ignore:
- 'docs/**'
- '**/*.md'
pull_request:
branches: [ main ]
paths-ignore:
- 'docs/**'
- '**/*.md'

permissions:
contents: read
id-token: write # required for OIDC authentication with CodSpeed

jobs:
benchmarks:
name: Benchmarks
runs-on: ubuntu-latest
services:
ministack:
image: ministackorg/ministack:1.3.53
ports:
- 4566:4566
env:
AWS_DEFAULT_REGION: us-east-1
GATEWAY_PORT: 4566
MINISTACK_ACCOUNT_ID: "000000000000"
MINISTACK_REGION: us-east-1
LOG_LEVEL: INFO
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
persist-credentials: false
- uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.13"
- uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2
with:
enable-cache: true
version: "latest"
- name: Install dependencies
run: uv sync --all-extras
- name: Run the benchmarks
uses: CodSpeedHQ/action@4296e51e7041e24dadb86d1d6e8b9320d223dbe8 # v5.0.3
with:
mode: simulation
run: uv run pytest tests/ --codspeed
2 changes: 1 addition & 1 deletion docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ services:
LOG_LEVEL: INFO
PERSIST_STATE: "1"
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:4566/_ministack/health"]
test: ["CMD", "busybox", "wget", "-q", "-O", "/dev/null", "http://127.0.0.1:4566/_ministack/health"]
interval: 2s
timeout: 5s
retries: 30
Expand Down
11 changes: 7 additions & 4 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -44,19 +44,20 @@ dev = [
{include-group = "lint"},
{include-group = "test"},
{include-group = "types"},
"prek>=0.4.3",
"prek>=0.5.0",
]
test = [
"pytest>=9.0.3",
"pytest-asyncio>=0.23.8",
"pytest-codspeed>=5.0.3",
]
lint = [
"bandit>=1.9.4",
"ruff>=0.15.15",
"zizmor>=1.25.2",
"ruff>=0.16.5",
"zizmor>=1.29.0",
]
types = [
"mypy>=2.1.0",
"mypy>=2.3.1",
"types-aiobotocore[essential]>=3.7.0",
]
examples = [
Expand All @@ -82,6 +83,7 @@ asyncio_mode = "auto"
markers = [
"unit: marks unit tests",
"integration: marks tests with real infrastructure env",
"benchmark: marks performance benchmarks (run with --codspeed)",
]

[tool.coverage.report]
Expand Down Expand Up @@ -129,6 +131,7 @@ ignore = [

"EM101",
"TRY003",
"CPY001",
]

[tool.ruff.lint.per-file-ignores]
Expand Down
Empty file added tests/benchmarks/__init__.py
Empty file.
62 changes: 62 additions & 0 deletions tests/benchmarks/conftest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
from collections.abc import AsyncGenerator

import pytest
from taskiq.acks import AckableMessage
from taskiq.result import TaskiqResult

from tests.conftest import AWSCredentials, _queue_name_from_url

from taskiq_sqs import S3ResultBackend, SQSBroker


BATCH_SIZE = 10
RESULT_TASK_ID = "benchmark-task"


async def ack(message: AckableMessage) -> None:
result = message.ack()
if result is not None:
await result


@pytest.fixture
async def bench_broker(
aws_credentials: AWSCredentials,
sqs_queue: str,
) -> AsyncGenerator[SQSBroker, None]:
"""The shared `sqs_broker`, but tuned for measuring a round-trip.

Two settings differ from the default broker: long polling, because with `wait_time_seconds=0` a receive returns
immediately and usually empty, so a round-trip would measure idle polls; and a batch size that lets a single
`ReceiveMessage` call serve the whole batch.
"""
broker = SQSBroker(
queue_name=_queue_name_from_url(sqs_queue),
wait_time_seconds=1,
max_number_of_messages=BATCH_SIZE,
**aws_credentials,
)
await broker.startup()
yield broker
await broker.shutdown()


@pytest.fixture
async def listener(bench_broker: SQSBroker) -> AsyncGenerator[AsyncGenerator[AckableMessage, None], None]:
"""A single `listen()` generator, created outside the measured test body."""
messages = bench_broker.listen()
yield messages
await messages.aclose()


@pytest.fixture
async def stored_result(s3_backend: S3ResultBackend[str]) -> TaskiqResult[str]:
"""A result already written to the bucket, so read benchmarks only measure reads."""
result: TaskiqResult[str] = TaskiqResult(
is_err=False,
return_value="benchmark",
execution_time=0.1,
log=None,
)
await s3_backend.set_result(RESULT_TASK_ID, result)
return result
43 changes: 43 additions & 0 deletions tests/benchmarks/test_broker.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
from collections.abc import AsyncGenerator

import pytest
from taskiq import BrokerMessage
from taskiq.acks import AckableMessage

from .conftest import BATCH_SIZE, ack
from taskiq_sqs import SQSBroker


@pytest.mark.benchmark
async def test_build_kick_kwargs(bench_broker: SQSBroker, broker_message: BrokerMessage) -> None:
await bench_broker._build_kick_kwargs(broker_message)


@pytest.mark.benchmark
async def test_kick(bench_broker: SQSBroker, broker_message: BrokerMessage) -> None:
await bench_broker.kick(broker_message)


@pytest.mark.benchmark
async def test_kick_and_listen_roundtrip(
bench_broker: SQSBroker,
broker_message: BrokerMessage,
listener: AsyncGenerator[AckableMessage, None],
) -> None:
await bench_broker.kick(broker_message)
message = await anext(listener)
await ack(message)


@pytest.mark.benchmark
async def test_kick_and_listen_roundtrip_batch(
bench_broker: SQSBroker,
broker_message: BrokerMessage,
listener: AsyncGenerator[AckableMessage, None],
) -> None:
for _ in range(BATCH_SIZE):
await bench_broker.kick(broker_message)
# SQS may return fewer messages than requested per call, so keep pulling until the whole batch is back.
for _ in range(BATCH_SIZE):
message = await anext(listener)
await ack(message)
27 changes: 27 additions & 0 deletions tests/benchmarks/test_result_backend.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
import pytest
from taskiq.result import TaskiqResult

from tests.benchmarks.conftest import RESULT_TASK_ID

from taskiq_sqs import S3ResultBackend


@pytest.mark.benchmark
async def test_set_result(s3_backend: S3ResultBackend[str], stored_result: TaskiqResult[str]) -> None:
await s3_backend.set_result(RESULT_TASK_ID, stored_result)


@pytest.mark.benchmark
async def test_get_result(
s3_backend: S3ResultBackend[str],
stored_result: TaskiqResult[str], # noqa: ARG001
) -> None:
await s3_backend.get_result(RESULT_TASK_ID)


@pytest.mark.benchmark
async def test_is_result_ready(
s3_backend: S3ResultBackend[str],
stored_result: TaskiqResult[str], # noqa: ARG001
) -> None:
await s3_backend.is_result_ready(RESULT_TASK_ID)
Loading