Skip to content

Commit 9f37908

Browse files
feat(coverage): Add coverage.py code coverage collection (#87)
1 parent f605362 commit 9f37908

11 files changed

Lines changed: 1213 additions & 61 deletions

File tree

docs/coverage.md

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
# Code Coverage (Python)
2+
3+
The Python SDK collects per-test code coverage during Tusk Drift replay using `coverage.py`. Unlike Node.js (which uses V8's built-in coverage), Python requires the `coverage` package to be installed.
4+
5+
## Requirements
6+
7+
```bash
8+
pip install tusk-drift-python-sdk[coverage]
9+
```
10+
11+
If `coverage` is not installed when coverage is enabled, the SDK logs a warning and coverage is skipped. Tests still run normally.
12+
13+
## How It Works
14+
15+
### coverage.py Integration
16+
17+
When coverage is enabled (via `--show-coverage`, `--coverage-output`, or `coverage.enabled: true` in config), the CLI sets `TUSK_COVERAGE=true`. The SDK detects this during initialization and starts coverage.py:
18+
19+
```python
20+
# What the SDK does internally:
21+
import coverage
22+
cov = coverage.Coverage(
23+
source=[os.path.realpath(os.getcwd())],
24+
branch=True,
25+
omit=["*/site-packages/*", "*/venv/*", "*/.venv/*", "*/tests/*", "*/test_*.py", "*/__pycache__/*"],
26+
)
27+
cov.start()
28+
```
29+
30+
Key points:
31+
- `branch=True` enables branch coverage (arc-based tracking)
32+
- `source` is set to the real path of the working directory (symlinks resolved)
33+
- Third-party code (site-packages, venv) is excluded by default
34+
35+
### Snapshot Flow
36+
37+
1. **Baseline**: CLI sends `CoverageSnapshotRequest(baseline=true)`. The SDK:
38+
- Calls `cov.stop()`
39+
- Uses `cov.analysis2(filename)` for each measured file to get ALL coverable lines (statements + missing)
40+
- Returns lines with count=0 for uncovered, count=1 for covered
41+
- Calls `cov.erase()` then `cov.start()` to reset counters
42+
43+
2. **Per-test**: CLI sends `CoverageSnapshotRequest(baseline=false)`. The SDK:
44+
- Calls `cov.stop()`
45+
- Uses `cov.get_data().lines(filename)` to get only executed lines since last reset
46+
- Returns only covered lines (count=1)
47+
- Calls `cov.erase()` then `cov.start()` to reset
48+
49+
3. **Communication**: Results are sent back to the CLI via the existing protobuf channel — same socket used for replay. No HTTP server or extra ports.
50+
51+
### Branch Coverage
52+
53+
Branch coverage uses coverage.py's arc tracking. The SDK extracts per-line branch data using:
54+
55+
```python
56+
analysis = cov._analyze(filename) # Private API
57+
missing_arcs = analysis.missing_branch_arcs()
58+
executed_arcs = set(data.arcs(filename) or [])
59+
```
60+
61+
For each branch point (line with multiple execution paths), the SDK reports:
62+
- `total`: number of branch paths from that line
63+
- `covered`: number of paths that were actually taken
64+
65+
**Note:** `_analyze()` is a private coverage.py API. It's the only way to get per-line branch arc data. The public API (`analysis2()`) only provides aggregate branch counts. This means branch coverage may break on major coverage.py version upgrades.
66+
67+
### Path Handling
68+
69+
The SDK uses `os.path.realpath()` for the source root to handle symlinked project directories. File paths reported by coverage.py are also resolved via `realpath` before comparison. This prevents the silent failure where all files get filtered out because symlink paths don't match.
70+
71+
## Environment Variables
72+
73+
Set automatically by the CLI. You should not set these manually.
74+
75+
| Variable | Description |
76+
|----------|-------------|
77+
| `TUSK_COVERAGE` | Set to `true` by the CLI when coverage is enabled. The SDK checks this to decide whether to start coverage.py. |
78+
79+
Note: `NODE_V8_COVERAGE` is also set by the CLI (for Node.js), but the Python SDK ignores it — it only checks `TUSK_COVERAGE`.
80+
81+
## Thread Safety
82+
83+
Coverage collection uses a module-level lock (`threading.Lock`) to ensure thread safety:
84+
85+
- `start_coverage_collection()`: Acquires lock while initializing. Guards against double initialization — if called twice, stops the existing instance first.
86+
- `take_coverage_snapshot()`: Acquires lock for the entire stop/read/erase/start cycle.
87+
- `stop_coverage_collection()`: Acquires lock while stopping and cleaning up.
88+
89+
This is important because the protobuf communicator runs coverage handlers in a background thread.
90+
91+
## Limitations
92+
93+
- **`coverage` package required**: Unlike Node.js (V8 coverage is built-in), Python needs `pip install coverage`. If not installed, coverage silently doesn't work (warning logged).
94+
- **Performance overhead**: coverage.py uses `sys.settrace()` which adds 10-30% execution overhead. This only applies during coverage replay runs.
95+
- **Multi-process servers**: gunicorn with `--workers > 1` forks worker processes. The SDK starts coverage.py in the main process; forked workers don't inherit it. Use `--workers 1` during coverage runs.
96+
- **Private API for branches**: `_analyze()` is not part of coverage.py's public API. Branch coverage detail may break on future coverage.py versions.
97+
- **Python 3.12+ recommended for async**: coverage.py's `sys.settrace` can miss some async lines on Python < 3.12. Python 3.12+ uses `sys.monitoring` for better async tracking.
98+
- **Startup ordering**: coverage.py starts during SDK initialization. Code that executes before `TuskDrift.initialize()` (e.g., module-level code in `tusk_drift_init.py`) isn't tracked. This is why `tusk_drift_init.py` typically shows 0% coverage.
99+
- **C extensions invisible**: coverage.py can't track C extensions (numpy, Cython modules). Not relevant for typical web API servers.

docs/environment-variables.md

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,20 @@ These variables configure how the SDK connects to the Tusk CLI during replay:
174174

175175
These are typically set automatically by the Tusk CLI and do not need to be configured manually.
176176

177+
## Coverage Variables
178+
179+
Set automatically by the CLI when `tusk drift run --coverage` is used. You should **not** set them manually.
180+
181+
| Variable | Description |
182+
|----------|-------------|
183+
| `TUSK_COVERAGE` | Set to `true` when coverage is enabled. The SDK checks this to start coverage.py. |
184+
185+
Note: `NODE_V8_COVERAGE` is also set by the CLI (for Node.js) but is ignored by the Python SDK.
186+
187+
See [Coverage Guide](./coverage.md) for details on how coverage collection works.
188+
177189
## Related Docs
178190

179191
- [Initialization Guide](./initialization.md) - SDK initialization parameters and config file settings
180192
- [Quick Start Guide](./quickstart.md) - Record and replay your first trace
193+
- [Coverage Guide](./coverage.md) - Code coverage during test replay

drift/core/communication/__init__.py

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,38 +6,25 @@
66

77
from .communicator import CommunicatorConfig, ProtobufCommunicator
88
from .types import (
9-
CliMessage,
10-
CLIMessageType,
119
ConnectRequest,
1210
ConnectResponse,
1311
GetMockRequest,
1412
GetMockResponse,
15-
MessageType,
1613
MockRequestInput,
1714
MockResponseOutput,
18-
# Protobuf types (re-exported)
19-
SdkMessage,
20-
SDKMessageType,
2115
dict_to_span,
2216
extract_response_data,
2317
span_to_proto,
2418
)
2519

2620
__all__ = [
27-
# Message types
28-
"MessageType",
29-
"SDKMessageType",
30-
"CLIMessageType",
3121
# Request/Response types
3222
"ConnectRequest",
3323
"ConnectResponse",
3424
"GetMockRequest",
3525
"GetMockResponse",
3626
"MockRequestInput",
3727
"MockResponseOutput",
38-
# Protobuf types
39-
"SdkMessage",
40-
"CliMessage",
4128
# Utilities
4229
"span_to_proto",
4330
"dict_to_span",

drift/core/communication/communicator.py

Lines changed: 73 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,24 +10,31 @@
1010
from dataclasses import dataclass
1111
from typing import Any
1212

13-
from tusk.drift.core.v1 import GetMockRequest as ProtoGetMockRequest
13+
from tusk.drift.core.v1 import (
14+
BranchInfo,
15+
CliMessage,
16+
CoverageSnapshotResponse,
17+
FileCoverageData,
18+
InstrumentationVersionMismatchAlert,
19+
MessageType,
20+
SdkMessage,
21+
SendAlertRequest,
22+
SendInboundSpanForReplayRequest,
23+
SetTimeTravelResponse,
24+
UnpatchedDependencyAlert,
25+
)
26+
from tusk.drift.core.v1 import (
27+
GetMockRequest as ProtoGetMockRequest,
28+
)
1429

1530
from ...version import MIN_CLI_VERSION, SDK_VERSION
1631
from ..span_serialization import clean_span_to_proto
1732
from ..types import CleanSpanData, calling_library_context
1833
from .types import (
19-
CliMessage,
2034
ConnectRequest,
2135
GetMockRequest,
22-
InstrumentationVersionMismatchAlert,
23-
MessageType,
2436
MockRequestInput,
2537
MockResponseOutput,
26-
SdkMessage,
27-
SendAlertRequest,
28-
SendInboundSpanForReplayRequest,
29-
SetTimeTravelResponse,
30-
UnpatchedDependencyAlert,
3138
span_to_proto,
3239
)
3340

@@ -750,6 +757,10 @@ def _background_read_loop(self) -> None:
750757
self._handle_set_time_travel_sync(cli_message)
751758
continue
752759

760+
if cli_message.type == MessageType.COVERAGE_SNAPSHOT:
761+
self._handle_coverage_snapshot_sync(cli_message)
762+
continue
763+
753764
# Route responses to waiting callers by request_id
754765
request_id = cli_message.request_id
755766
if request_id:
@@ -774,8 +785,8 @@ def _background_read_loop(self) -> None:
774785

775786
def _handle_set_time_travel_sync(self, cli_message: CliMessage) -> None:
776787
"""Handle SetTimeTravel request from CLI and send response."""
777-
request = cli_message.set_time_travel_request
778-
if not request:
788+
request = getattr(cli_message, "set_time_travel_request", None)
789+
if request is None:
779790
return
780791

781792
logger.debug(
@@ -809,6 +820,57 @@ def _handle_set_time_travel_sync(self, cli_message: CliMessage) -> None:
809820
except Exception as e:
810821
logger.error(f"Failed to send SetTimeTravel response: {e}")
811822

823+
def _handle_coverage_snapshot_sync(self, cli_message: CliMessage) -> None:
824+
"""Handle CoverageSnapshot request from CLI and send response."""
825+
request = getattr(cli_message, "coverage_snapshot_request", None)
826+
if request is None:
827+
return
828+
829+
logger.debug(f"Received CoverageSnapshot request: baseline={request.baseline}")
830+
831+
try:
832+
from ..coverage_server import take_coverage_snapshot
833+
834+
result = take_coverage_snapshot(request.baseline)
835+
836+
# Convert to protobuf
837+
coverage: dict[str, FileCoverageData] = {}
838+
for file_path, file_data in result.items():
839+
branches: dict[str, BranchInfo] = {}
840+
for line, branch_info in file_data.get("branches", {}).items():
841+
branches[line] = BranchInfo(
842+
total=branch_info.get("total", 0),
843+
covered=branch_info.get("covered", 0),
844+
)
845+
846+
coverage[file_path] = FileCoverageData(
847+
lines=file_data.get("lines", {}),
848+
total_branches=file_data.get("totalBranches", 0),
849+
covered_branches=file_data.get("coveredBranches", 0),
850+
branches=branches,
851+
)
852+
853+
response = CoverageSnapshotResponse(
854+
success=True,
855+
error="",
856+
coverage=coverage,
857+
)
858+
except Exception as e:
859+
logger.error(f"Failed to take coverage snapshot: {e}")
860+
response = CoverageSnapshotResponse(success=False, error=str(e))
861+
862+
sdk_message = SdkMessage(
863+
type=MessageType.COVERAGE_SNAPSHOT,
864+
request_id=cli_message.request_id,
865+
coverage_snapshot_response=response,
866+
)
867+
868+
try:
869+
self._send_message_sync(sdk_message)
870+
logger.debug(f"Sent CoverageSnapshot response: success={response.success}")
871+
except Exception as e:
872+
logger.error(f"[coverage] Failed to send response: {e}")
873+
812874
def _send_message_sync(self, message: SdkMessage) -> None:
813875
"""Send a message synchronously on the main socket."""
814876
if not self._socket:

drift/core/communication/types.py

Lines changed: 3 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -12,20 +12,6 @@
1212
from __future__ import annotations
1313

1414
__all__ = [
15-
# Re-exported protobuf types
16-
"CliMessage",
17-
"InstrumentationVersionMismatchAlert",
18-
"MessageType",
19-
"Runtime",
20-
"SdkMessage",
21-
"SendAlertRequest",
22-
"SendInboundSpanForReplayRequest",
23-
"SetTimeTravelRequest",
24-
"SetTimeTravelResponse",
25-
"UnpatchedDependencyAlert",
26-
# Aliases
27-
"SDKMessageType",
28-
"CLIMessageType",
2915
# Dataclasses
3016
"ConnectRequest",
3117
"ConnectResponse",
@@ -42,18 +28,6 @@
4228
from dataclasses import dataclass, field
4329
from typing import Any
4430

45-
from tusk.drift.core.v1 import (
46-
CliMessage,
47-
InstrumentationVersionMismatchAlert,
48-
MessageType,
49-
Runtime,
50-
SdkMessage,
51-
SendAlertRequest,
52-
SendInboundSpanForReplayRequest,
53-
SetTimeTravelRequest,
54-
SetTimeTravelResponse,
55-
UnpatchedDependencyAlert,
56-
)
5731
from tusk.drift.core.v1 import (
5832
ConnectRequest as ProtoConnectRequest,
5933
)
@@ -66,6 +40,9 @@
6640
from tusk.drift.core.v1 import (
6741
GetMockResponse as ProtoGetMockResponse,
6842
)
43+
from tusk.drift.core.v1 import (
44+
Runtime,
45+
)
6946
from tusk.drift.core.v1 import (
7047
Span as ProtoSpan,
7148
)
@@ -79,9 +56,6 @@
7956
StatusCode as ProtoStatusCode,
8057
)
8158

82-
SDKMessageType = MessageType
83-
CLIMessageType = MessageType
84-
8559

8660
def _python_to_value(value: Any) -> Any:
8761
"""Convert Python value to protobuf Value."""

0 commit comments

Comments
 (0)