feat: add stress tests and CI benchmark integration

.github/workflows/ci-pydgraph-benchmarks.yml

@@ -0,0 +1,34 @@
+name: ci-pydgraph-benchmarks
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - v[0-9]+.[0-9]+.[0-9]+*
+
+permissions:
+  contents: read
+
+jobs:
+  benchmarks:
+    name: Release Benchmarks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+      - name: Setup python runtime and tooling
+        uses: ./.github/actions/setup-python-and-tooling
+        with:
+          python-version: "3.13"
+      - name: Setup project dependencies
+        run: INSTALL_MISSING_TOOLS=true make setup
+      - name: Sync python virtualenv
+        run: make sync
+      - name: Run benchmarks
+        run: make benchmark
+      - name: Upload benchmark results
+        uses: actions/upload-artifact@v4
+        with:
+          name: benchmark-results-${{ github.ref_name }}
+          path: |
+            benchmark-results.json
+            benchmark-histogram.svg

.gitignore

@@ -41,3 +41,16 @@ venv
 pyvenv.cfg
 .DS_Store
 examples/notebook/RAG/.env
+.osgrep
+
+# Git worktrees
+.worktrees/
+
+# Benchmark outputs
+benchmark-results.json
+benchmark-histogram.svg
+stress-benchmark-results.json
+
+# Downloaded test data (fetched on demand)
+tests/resources/1million.rdf.gz
+tests/resources/1million.schema

.pre-commit-config.yaml

@@ -65,6 +65,7 @@ repos:
         pass_filenames: false
         additional_dependencies:
           - pytest>=8.3.3
+          - pytest-benchmark>=4.0.0
           - grpcio>=1.65.1
           - protobuf>=4.23.0
   - repo: https://github.com/pre-commit/mirrors-mypy

CONTRIBUTING.md

@@ -220,15 +220,48 @@ make test
 Run specific tests:
 
 ```sh
-bash scripts/local-test.sh -v tests/test_connect.py::TestOpen
+make test PYTEST_ARGS="-v tests/test_connect.py::TestOpen"
 ```
 
 Run a single test:
 
 ```sh
-bash scripts/local-test.sh -v tests/test_connect.py::TestOpen::test_connection_with_auth
+make test PYTEST_ARGS="-v tests/test_connect.py::TestOpen::test_connection_with_auth"
 ```
 
+### Stress Tests
+
+The project includes comprehensive stress tests that verify concurrent operations, transaction
+conflicts, deadlock prevention, and retry mechanisms for both sync and async clients.
+
+**Quick mode** (default, ~12 seconds) - 20 workers, 50 ops, 10 iterations:
+
+```sh
+make test PYTEST_ARGS="tests/test_stress_sync.py tests/test_stress_async.py -v"
+```
+
+**Moderate mode** (10x quick, includes movie dataset, ~60+ seconds) - 200 workers, 500 ops, 100
+iterations:
+
+```sh
+make test STRESS_TEST_MODE=moderate PYTEST_ARGS="tests/test_stress_sync.py tests/test_stress_async.py -v"
+```
+
+**Full mode** (10x moderate, maximum stress, ~10+ minutes) - 2000 workers, 5000 ops, 1000
+iterations:
+
+```sh
+make test STRESS_TEST_MODE=full PYTEST_ARGS="tests/test_stress_sync.py tests/test_stress_async.py -v"
+```
+
+The stress tests cover:
+
+- **Sync tests**: Run with `ThreadPoolExecutor` to test concurrent operations
+- **Async tests**: Use pure `asyncio.gather()` concurrency (no `concurrent.futures` mixing)
+- **Retry utilities**: Tests for `retry_async()`, `with_retry_async()`, and
+  `run_transaction_async()`
+- **Deadlock regression**: Validates the asyncio.Lock deadlock fix from PR #293
+
 ### Test Infrastructure
 
 The test script requires Docker and Docker Compose to be installed on your machine.

Makefile

@@ -2,6 +2,25 @@
 SHELL := /bin/bash
 export PATH := $(HOME)/.local/bin:$(HOME)/.cargo/bin:$(PATH)
 
+# Export test configuration variables so they're available to child processes
+# Usage: make test STRESS_TEST_MODE=moderate PYTEST_ARGS="-v"
+#        make test LOG=info   (adds --log-cli-level=INFO to default PYTEST_ARGS)
+export STRESS_TEST_MODE
+export DGRAPH_IMAGE_TAG
+
+# When LOG is set (e.g., LOG=info), inject --log-cli-level into pytest flags.
+# Works with both the default PYTEST_ARGS and explicit overrides:
+#   make test LOG=info                      → -v --benchmark-disable --log-cli-level=INFO
+#   make benchmark LOG=warning              → --benchmark-only ... --log-cli-level=WARNING
+#   make test PYTEST_ARGS="-x" LOG=debug    → -x --log-cli-level=DEBUG
+PYTEST_ARGS ?= -v --benchmark-disable
+ifdef LOG
+  LOG_FLAG := --log-cli-level=$(shell echo '$(LOG)' | tr '[:lower:]' '[:upper:]')
+  PYTEST_ARGS += $(LOG_FLAG)
+endif
+export LOG
+export PYTEST_ARGS
+
 # Source venv if it exists and isn't already active
 PROJECT_VENV := $(CURDIR)/.venv
 ACTIVATE := $(wildcard .venv/bin/activate)
@@ -15,14 +34,21 @@ else
   RUN :=
 endif
 
-.PHONY: help setup sync deps deps-uv deps-trunk deps-docker test check protogen clean build publish
+.PHONY: help setup sync deps deps-uv deps-trunk deps-docker test benchmark check protogen clean build publish
 
 .DEFAULT_GOAL := help
 
 help: ## Show this help message
 	@echo ""
 	@echo "Environment Variables:"
 	@echo "  INSTALL_MISSING_TOOLS=true    Enable automatic installation of missing tools (default: disabled)"
+	@echo "  LOG=<level>                   Add --log-cli-level to pytest (e.g., LOG=info, LOG=debug)"
+	@echo "                                Works with both 'test' and 'benchmark' targets"
+	@echo "  STRESS_TEST_MODE=<mode>       Stress test preset: quick (default), moderate, full"
+	@echo "  PYTEST_ARGS=\"...\"             Override default pytest flags (default: -v --benchmark-disable)"
+	@echo "                                Note: overrides LOG when set explicitly. 'benchmark' sets its own"
+	@echo "                                PYTEST_ARGS internally but still honours LOG"
+	@echo "  DGRAPH_IMAGE_TAG=<tag>        Override the Dgraph Docker image tag (default: latest)"
 	@echo ""
 	@echo "Available targets:"
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-15s\033[0m %s\n", $$1, $$2}'
@@ -51,8 +77,31 @@ clean: ## Cleans build artifacts
 build: deps-uv sync protogen ## Builds release package
 	$(RUN) uv build
 
-test: deps-uv sync ## Run tests
-	bash scripts/local-test.sh
+test: deps-uv sync ## Run tests (use PYTEST_ARGS to pass options, e.g., make test PYTEST_ARGS="-v tests/test_connect.py")
+	bash scripts/local-test.sh $(PYTEST_ARGS)
+
+benchmark: ## Run benchmarks (measures per-operation latency with pytest-benchmark)
+	@# Outputs (all .gitignored):
+	@#   benchmark-results.json        Phase 1 results (pytest-benchmark JSON)
+	@#   benchmark-histogram.svg       Phase 1 latency histogram
+	@#   stress-benchmark-results.json Phase 2 results (pytest-benchmark JSON)
+	@#
+	@# Phase 1: Per-operation latency benchmarks against a clean database.
+	@# Runs targeted benchmark tests (test_benchmark_*.py) which measure individual
+	@# operations (query, mutation, upsert, etc.) in isolation.  Each test creates a
+	@# fresh schema via drop_all, so these MUST run on their own Dgraph cluster —
+	@# the rapid schema churn destabilises the alpha for any tests that follow.
+	@echo "═══ Phase 1: Per-operation latency benchmarks ═══"
+	$(MAKE) test PYTEST_ARGS="--benchmark-only --benchmark-json=benchmark-results.json --benchmark-histogram=benchmark-histogram -v $(LOG_FLAG) tests/test_benchmark_async.py tests/test_benchmark_sync.py"
+	@# Phase 2: Stress-test benchmarks under sustained concurrent load.
+	@# Runs stress tests (test_stress_*.py) with the 1-million-movie dataset loaded.
+	@# Uses a separate Dgraph cluster (via a second 'make test' invocation) so the
+	@# alpha starts fresh after Phase 1's drop_all churn.
+	@# benchmark.pedantic(rounds=1) in each stress test prevents pytest-benchmark
+	@# from compounding iterations — the stress_config["rounds"] inner loop
+	@# (controlled by STRESS_TEST_MODE) handles repetition instead.
+	@echo "═══ Phase 2: Stress-test benchmarks (moderate load, 1M movies) ═══"
+	$(MAKE) test STRESS_TEST_MODE=moderate PYTEST_ARGS="--benchmark-only --benchmark-json=stress-benchmark-results.json -v $(LOG_FLAG) tests/test_stress_async.py tests/test_stress_sync.py"
 
 publish: clean build  ## Publish a new release to PyPi (requires UV_PUBLISH_USERNAME and UV_PUBLISH_PASSWORD to be set)
 	$(RUN) uv publish

PUBLISHING.md

@@ -9,12 +9,12 @@ This document contains instructions to create a new pydgraph release and publish
 1. Create a new branch (prepare-for-release-vXX.X.X, for instance)
 1. Update the VERSION in pydgraph/meta.py
 1. Build pydgraph locally, see the [README](README.md#build-from-source)
-1. Run the tests (`bash scripts/local-test.sh`) to ensure everything works
+1. Run the tests (`make test`) to ensure everything works
 1. If you're concerned about incompatibilities with earlier Dgraph versions, invoke the test suite
    with earlier Dgraph versions
 
    ```sh
-   DGRAPH_IMAGE_TAG=vX.X.X bash scripts/local-test.sh
+   make test DGRAPH_IMAGE_TAG=vX.X.X
    ```
 
 1. If you happen to have the testpypi access token, try a test upload to testpypi: