Skip to content

Commit 1d785d9

Browse files
jy-tanjy-tan
authored
feat: grpc client instrumentation (#47)
1 parent 16232cc commit 1d785d9

21 files changed

Lines changed: 3530 additions & 2 deletions

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ Tusk Drift currently supports the following packages and versions:
5252
| HTTPX | all versions |
5353
| aiohttp | all versions |
5454
| urllib3 | all versions |
55+
| grpcio (client-side only) | all versions |
5556
| psycopg | `>=3.1.12` |
5657
| psycopg2 | all versions |
5758
| Redis | `>=4.0.0` |

‎drift/core/drift_sdk.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,16 @@ def _init_auto_instrumentations(self) -> None:
451451
except ImportError:
452452
pass
453453

454+
try:
455+
import grpc # type: ignore[unresolved-import]
456+
457+
from ..instrumentation.grpc import GrpcInstrumentation
458+
459+
_ = GrpcInstrumentation()
460+
logger.debug("gRPC instrumentation initialized")
461+
except ImportError:
462+
pass
463+
454464
try:
455465
import django
456466

‎drift/instrumentation/grpc/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
"""gRPC client instrumentation."""
2+
3+
from .instrumentation import GrpcInstrumentation
4+
5+
__all__ = ["GrpcInstrumentation"]

‎drift/instrumentation/grpc/e2e-tests/.tusk/config.yaml‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
version: 1
2+
3+
service:
4+
id: "grpc-e2e-test-id"
5+
name: "grpc-e2e-test"
6+
port: 8000
7+
start:
8+
command: "python src/app.py"
9+
readiness_check:
10+
command: "curl -f http://localhost:8000/health"
11+
timeout: 45s
12+
interval: 5s
13+
14+
tusk_api:
15+
url: "http://localhost:8000"
16+
17+
test_execution:
18+
concurrent_limit: 10
19+
batch_size: 10
20+
timeout: 30s
21+
22+
recording:
23+
sampling_rate: 1.0
24+
export_spans: false
25+
26+
replay:
27+
enable_telemetry: false

‎drift/instrumentation/grpc/e2e-tests/Dockerfile‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
FROM python-e2e-base:latest
2+
3+
# Copy SDK source for editable install
4+
COPY . /sdk
5+
6+
# Copy test files
7+
COPY drift/instrumentation/grpc/e2e-tests /app
8+
9+
WORKDIR /app
10+
11+
# Install dependencies (requirements.txt uses -e /sdk for SDK)
12+
RUN pip install -q -r requirements.txt
13+
14+
# Make entrypoint executable
15+
RUN chmod +x entrypoint.py
16+
17+
# Create .tusk directories
18+
RUN mkdir -p /app/.tusk/traces /app/.tusk/logs
19+
20+
# Run entrypoint
21+
ENTRYPOINT ["python", "entrypoint.py"]

‎drift/instrumentation/grpc/e2e-tests/docker-compose.yml‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
services:
2+
app:
3+
build:
4+
context: ../../../..
5+
dockerfile: drift/instrumentation/grpc/e2e-tests/Dockerfile
6+
args:
7+
- TUSK_CLI_VERSION=${TUSK_CLI_VERSION:-latest}
8+
environment:
9+
- PORT=8000
10+
- TUSK_ANALYTICS_DISABLED=1
11+
- PYTHONUNBUFFERED=1
12+
working_dir: /app
13+
volumes:
14+
# Mount SDK source for hot reload (no rebuild needed for SDK changes)
15+
- ../../../..:/sdk
16+
# Mount app source for development
17+
- ./src:/app/src
18+
# Mount .tusk folder to persist traces
19+
- ./.tusk:/app/.tusk

‎drift/instrumentation/grpc/e2e-tests/entrypoint.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
#!/usr/bin/env python3
2+
"""
3+
E2E Test Entrypoint for gRPC Instrumentation
4+
5+
This script orchestrates the full e2e test lifecycle:
6+
1. Setup: Install dependencies, generate proto files
7+
2. Record: Start app in RECORD mode, execute requests
8+
3. Test: Run Tusk CLI tests
9+
4. Teardown: Cleanup and return exit code
10+
"""
11+
12+
import sys
13+
from pathlib import Path
14+
15+
# Add SDK to path for imports
16+
sys.path.insert(0, "/sdk")
17+
18+
from drift.instrumentation.e2e_common.base_runner import E2ETestRunnerBase
19+
20+
21+
class GrpcE2ETestRunner(E2ETestRunnerBase):
22+
"""E2E test runner for gRPC instrumentation."""
23+
24+
def __init__(self):
25+
import os
26+
27+
port = int(os.getenv("PORT", "8000"))
28+
super().__init__(app_port=port)
29+
30+
def setup(self):
31+
"""Phase 1: Setup dependencies and generate proto files."""
32+
self.log("=" * 50, self.Colors.BLUE)
33+
self.log("Phase 1: Setup", self.Colors.BLUE)
34+
self.log("=" * 50, self.Colors.BLUE)
35+
36+
self.log("Installing Python dependencies...", self.Colors.BLUE)
37+
self.run_command(["pip", "install", "-q", "-r", "requirements.txt"])
38+
39+
# Generate proto files
40+
self.log("Generating proto files...", self.Colors.BLUE)
41+
self.run_command(
42+
[
43+
"python",
44+
"-m",
45+
"grpc_tools.protoc",
46+
"-I",
47+
"src/proto",
48+
"--python_out=src",
49+
"--grpc_python_out=src",
50+
"src/proto/greeter.proto",
51+
]
52+
)
53+
54+
self.log("Setup complete", self.Colors.GREEN)
55+
56+
# Use Colors from base class
57+
@property
58+
def Colors(self):
59+
from drift.instrumentation.e2e_common.base_runner import Colors
60+
61+
return Colors
62+
63+
64+
if __name__ == "__main__":
65+
runner = GrpcE2ETestRunner()
66+
exit_code = runner.run()
67+
sys.exit(exit_code)

‎drift/instrumentation/grpc/e2e-tests/requirements.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
-e /sdk
2+
Flask>=3.1.2
3+
grpcio>=1.60.0
4+
grpcio-tools>=1.60.0
5+
protobuf>=6.0

‎drift/instrumentation/grpc/e2e-tests/run.sh‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#!/bin/bash
2+
3+
# Exit on error
4+
set -e
5+
6+
# Accept optional port parameter (default: 8000)
7+
APP_PORT=${1:-8000}
8+
export APP_PORT
9+
10+
# Generate unique docker compose project name
11+
# Get the instrumentation name (parent directory of e2e-tests)
12+
TEST_NAME="$(basename "$(dirname "$(pwd)")")"
13+
PROJECT_NAME="python-${TEST_NAME}-${APP_PORT}"
14+
15+
# Colors for output
16+
GREEN='\033[0;32m'
17+
RED='\033[0;31m'
18+
YELLOW='\033[1;33m'
19+
BLUE='\033[0;34m'
20+
NC='\033[0m'
21+
22+
echo -e "${BLUE}========================================${NC}"
23+
echo -e "${BLUE}Running Python E2E Test: ${TEST_NAME}${NC}"
24+
echo -e "${BLUE}Port: ${APP_PORT}${NC}"
25+
echo -e "${BLUE}========================================${NC}"
26+
echo ""
27+
28+
# Cleanup function
29+
cleanup() {
30+
echo ""
31+
echo -e "${YELLOW}Cleaning up containers...${NC}"
32+
docker compose -p "$PROJECT_NAME" down -v 2>/dev/null || true
33+
}
34+
35+
# Register cleanup on exit
36+
trap cleanup EXIT
37+
38+
# Build containers
39+
echo -e "${BLUE}Building containers...${NC}"
40+
docker compose -p "$PROJECT_NAME" build --no-cache
41+
42+
# Run the test container
43+
echo -e "${BLUE}Starting test...${NC}"
44+
echo ""
45+
46+
# Run container and capture exit code (always use port 8000 inside container)
47+
# Disable set -e temporarily to capture exit code
48+
set +e
49+
docker compose -p "$PROJECT_NAME" run --rm app
50+
EXIT_CODE=$?
51+
set -e
52+
53+
echo ""
54+
if [ $EXIT_CODE -eq 0 ]; then
55+
echo -e "${GREEN}========================================${NC}"
56+
echo -e "${GREEN}Test passed!${NC}"
57+
echo -e "${GREEN}========================================${NC}"
58+
else
59+
echo -e "${RED}========================================${NC}"
60+
echo -e "${RED}Test failed with exit code ${EXIT_CODE}${NC}"
61+
echo -e "${RED}========================================${NC}"
62+
fi
63+
64+
exit $EXIT_CODE

0 commit comments

Comments
 (0)