feat(coverage): Add coverage.py code coverage collection by sohil-kshirsagar · Pull Request #87 · Use-Tusk/drift-python-sdk

sohil-kshirsagar · 2026-04-03T21:42:17Z

Add code coverage collection using coverage.py to the Python SDK. When the CLI enables coverage, the SDK manages coverage.py programmatically and returns per-file line/branch coverage via protobuf.

How it works

CLI sets TUSK_COVERAGE=true env var
SDK starts coverage.py with branch=True during initialization
CLI sends CoverageSnapshotRequest via protobuf channel
SDK: stop() → get_data() / analysis2() → erase() → start() cycle
SDK returns CoverageSnapshotResponse with per-file data

Key design decisions

SDK-driven coverage.py: Python has no equivalent of NODE_V8_COVERAGE. The SDK runs inside the app process and drives coverage.py directly.
Branch coverage via _analyze(): Uses coverage.py's private _analyze() API for arc-based branch tracking. The public API only provides aggregate branch counts. Documented as a known limitation.
Thread-safe: All coverage operations protected by threading.Lock(). Safe for concurrent protobuf handler access.

Files changed

drift/core/coverage_server.py — Coverage lifecycle management (new file)
drift/core/communication/communicator.py — Coverage snapshot handler + betterproto fixes
drift/core/communication/types.py — Import new proto types
drift/drift_sdk.py — Hook coverage start/stop into SDK lifecycle
tests/unit/test_coverage_server.py — Unit tests for coverage server
docs/coverage.md — SDK coverage internals documentation
docs/environment-variables.md — Coverage env vars section

Also includes (non-coverage)

betterproto getattr() fix: SetTimeTravelRequest handler used if not request: which fails for messages with all-default values (betterproto falsy bug). Fixed with getattr(cli_message, "field", None). Applied to both time travel and coverage handlers.

TODOs before merge

Blocked on feat(proto): Add coverage tracking messages tusk-drift-schemas#47 — tusk-drift-schemas updated to >=0.1.34
~~Add unit tests for coverage_server.py~~
~~Guard coverage start with REPLAY mode check (avoid overhead in RECORD/DISABLED modes)~~
Consider pinning minimum coverage.py version for _analyze() API stability

Edge cases / gotchas

_is_user_file uses os.sep trailing separator to prevent /app matching /application
os.path.realpath() resolves symlinks for consistent path comparison
Double-init guard: calling start_coverage_collection() twice stops the existing instance first
*/test* omit was too broad (matched testimony.py), narrowed to */tests/* and */test_*.py

When TUSK_COVERAGE_PORT env var is set, the SDK starts a tiny HTTP server that manages coverage.py. On each /snapshot request: - Stop coverage, get data, erase (reset), restart - Returns per-file line counts with clean per-test data - No diffing needed (coverage.py supports stop/erase/start cycle) Works with Flask, FastAPI, Django, gunicorn, uvicorn - any framework, because the SDK runs inside the app process. Requires: pip install coverage (or pip install tusk-drift[coverage])

When /snapshot?baseline=true is called, uses coverage.analysis2() to get ALL coverable statements (including uncovered) for the denominator. Regular /snapshot calls only return executed lines (for per-test data).

- Threading lock protects stop/get_data/erase/start sequence - stop_coverage_server() for clean shutdown, integrated into SDK shutdown() - Module-level server reference for proper cleanup

- Enable branch=True in coverage.py initialization - Extract branch data via cov._analyze(filename) API: - numbers.n_branches, n_missing_branches for totals - missing_branch_arcs() for per-line branch detail - Return branch data in /snapshot response alongside line coverage - Python shows accurate branch coverage (93.3% for demo app)

…- needs debug)

betterproto treats messages with all default values as falsy. CoverageSnapshotRequest(baseline=False) was falsy, causing per-test snapshots to be skipped. Changed 'if not request' to 'if request is None'. Also separated coverage initialization from HTTP server so coverage.py starts via start_coverage_collection() for the protobuf channel. Extracted take_coverage_snapshot() as reusable function.

- Remove HTTP server code (CoverageSnapshotHandler, start_coverage_server, _coverage_server global, HTTPServer import) - Replace with clean module-level state (_cov_instance, _source_root, _lock) - Extract _is_user_file() helper - stop_coverage_server() -> stop_coverage_collection() - Update module docstring to reflect protobuf-only architecture

Add _lock protection to stop_coverage_collection() to prevent race condition where shutdown sets _cov_instance=None while a snapshot is in progress on the background reader thread.

Add docs/coverage.md explaining coverage.py integration, branch coverage via arc tracking, thread safety, and limitations. Update environment-variables.md with coverage env vars section.

- Use getattr() for betterproto oneof field access (prevents AttributeError) - Fix _is_user_file path prefix collision (/app matching /application) - Add os.path.realpath() for symlink-safe path comparison - Add thread lock to start_coverage_collection() - Add double-init guard (stop existing instance before creating new) - Narrow */test* omit pattern to */tests/* and */test_*.py - Log failed file analysis at debug level instead of silent swallow

drift/core/coverage_server.py

drift/core/drift_sdk.py

- Move _cov_instance None check inside lock (TOCTOU race fix) - Fix branch counting to only include actual branch points, not all arcs

drift/core/coverage_server.py

- Add None guard for _cov_instance before ._analyze() to fix type checker error - Wrap stop/erase/start cycle in try/finally so coverage always restarts on error - Run ruff format on coverage_server.py and drift_sdk.py - Update uv.lock to match pyproject.toml dependency changes

drift/core/coverage_server.py

…rageSnapshotResponse, FileCoverageData)

drift/core/coverage_server.py

….drift.core.v1 types.py now only contains SDK wrapper dataclasses and utility functions. Proto types are imported directly where used. Removes SDKMessageType/CLIMessageType aliases.

…tests - Skip coverage collection when TUSK_DRIFT_MODE is set to non-REPLAY mode - Add 9 unit tests covering start/stop, mode gating, file filtering, error handling

cubic-dev-ai

2 issues found across 10 files

Prompt for AI agents (unresolved issues)


Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="drift/core/drift_sdk.py">

<violation number="1" location="drift/core/drift_sdk.py:167">
P2: Coverage collection is started before the `_initialized` early-return, so repeated `initialize()` calls are no longer side-effect free and will restart coverage unexpectedly.</violation>
</file>

<file name="docs/coverage.md">

<violation number="1" location="docs/coverage.md:10">
P2: The docs reference `tusk-drift[coverage]`, but this package extra is not defined in `pyproject.toml`, so the install command is invalid.</violation>
</file>

_{Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review, or fix all with cubic.}

drift/core/drift_sdk.py

docs/coverage.md

…der init - Add coverage>=7.0.0 as optional dependency [coverage] extra in pyproject.toml - Fix docs to reference correct package name tusk-drift-python-sdk[coverage] - Move start_coverage_collection() after _initialized guard to avoid wasteful re-invocation

…branch counts Branch detection via _analyze() depends on observed arcs, which vary with thread timing. Now the baseline snapshot caches branch point structure (totals per line), and per-test snapshots reuse that cache — only computing covered counts from the current test's arcs. This eliminates flaky branch totals (was 12/18 or 12/22 randomly, now consistently 4/10).

tusk-dev · 2026-04-07T03:15:02Z

Already incorporated tests

Tip

New to Tusk Unit Tests? Learn more here.

Avg +34% line coverage gain across 2 files

Source file	Line	Branch
`drift/core/communication/communicator.py`	53% (+8%)	0%
`drift/core/coverage_server.py`	93% (+60%)	0%

Coverage is calculated by running tests directly associated with each source file, learn more here.

View check history

Commit	Status	Output	Created (UTC)
`c71dd37`	Skipped due to new commit on branch	Output	Apr 7, 2026 2:39AM
`06b6c27`	Generated 31 tests - 31 passed	Tests	Apr 7, 2026 2:57AM
`e1fa8c0`	Generated 33 tests - 33 passed	Tests	Apr 7, 2026 3:09AM
`9b197f7`	Commit incorporated tests	Output	Apr 7, 2026 5:12AM
`1eb2e39`	Already incorporated tests	Output	Apr 7, 2026 8:56PM
`c0fac85`	Already incorporated tests	Output	Apr 7, 2026 9:17PM

…andler

…source_root on stop - start_coverage_collection() was before the _initialized check, causing repeated initialize() calls to stop/restart coverage and lose data - stop_coverage_collection() now resets _source_root for cleanup completeness

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 1eb2e39. Configure here.}

drift/core/coverage_server.py

sohil-kshirsagar added 11 commits April 1, 2026 00:45

feat: add ?baseline=true parameter using coverage.py analysis2

3261a40

When /snapshot?baseline=true is called, uses coverage.analysis2() to get ALL coverable statements (including uncovered) for the denominator. Regular /snapshot calls only return executed lines (for per-test data).

chore: add thread safety lock and clean shutdown for coverage server

2a54664

- Threading lock protects stop/get_data/erase/start sequence - stop_coverage_server() for clean shutdown, integrated into SDK shutdown() - Module-level server reference for proper cleanup

wip: migrate coverage to protobuf channel (Python handler timing out …

020e2af

…- needs debug)

fix: prod readiness - thread-safe coverage shutdown

5e4aa97

Add _lock protection to stop_coverage_collection() to prevent race condition where shutdown sets _cov_instance=None while a snapshot is in progress on the background reader thread.

feat: use TUSK_COVERAGE instead of NODE_V8_COVERAGE for Python

18e8349

docs: add code coverage documentation

5d438df

Add docs/coverage.md explaining coverage.py integration, branch coverage via arc tracking, thread safety, and limitations. Update environment-variables.md with coverage env vars section.

sohil-kshirsagar mentioned this pull request Apr 3, 2026

feat(coverage): Add code coverage collection and export Use-Tusk/tusk-cli#216

Merged

3 tasks

cursor bot reviewed Apr 3, 2026

View reviewed changes

drift/core/coverage_server.py Show resolved Hide resolved

drift/core/coverage_server.py Outdated Show resolved Hide resolved

drift/core/drift_sdk.py Outdated Show resolved Hide resolved

sohil-kshirsagar added 2 commits April 3, 2026 15:02

docs: clean up AI writing patterns in coverage doc

5b3354b

fix: address bugbot review feedback

4d156c4

- Move _cov_instance None check inside lock (TOCTOU race fix) - Fix branch counting to only include actual branch points, not all arcs

cursor bot reviewed Apr 3, 2026

View reviewed changes

drift/core/coverage_server.py Show resolved Hide resolved

sohil-kshirsagar added 2 commits April 6, 2026 17:08

chore: update tusk-drift-schemas to >=0.1.34

a9eeb5e

cursor bot reviewed Apr 7, 2026

View reviewed changes

drift/core/coverage_server.py Show resolved Hide resolved

sohil-kshirsagar added 2 commits April 6, 2026 18:05

fix: remove unused imports and simplify _is_user_file return

5cc5a07

fix: restore re-exported imports removed by mistake (BranchInfo, Cove…

97d824e

…rageSnapshotResponse, FileCoverageData)

cursor bot reviewed Apr 7, 2026

View reviewed changes

drift/core/coverage_server.py Show resolved Hide resolved

drift/core/coverage_server.py Show resolved Hide resolved

sohil-kshirsagar added 2 commits April 6, 2026 18:29

ref: remove proto re-exports from types.py, import directly from tusk…

01a0b3a

….drift.core.v1 types.py now only contains SDK wrapper dataclasses and utility functions. Proto types are imported directly where used. Removes SDKMessageType/CLIMessageType aliases.

fix: guard coverage with REPLAY mode check, add coverage_server unit …

c71dd37

…tests - Skip coverage collection when TUSK_DRIFT_MODE is set to non-REPLAY mode - Add 9 unit tests covering start/stop, mode gating, file filtering, error handling

sohil-kshirsagar changed the title ~~WIP feat(coverage): Add coverage.py code coverage collection~~ feat(coverage): Add coverage.py code coverage collection Apr 7, 2026

sohil-kshirsagar marked this pull request as ready for review April 7, 2026 02:39

cubic-dev-ai bot reviewed Apr 7, 2026

View reviewed changes

drift/core/drift_sdk.py Outdated Show resolved Hide resolved

docs/coverage.md Outdated Show resolved Hide resolved

sohil-kshirsagar added 2 commits April 6, 2026 19:57

test: add Tusk-generated tests for coverage server and communicator h…

9b197f7

…andler

sohil-kshirsagar requested a review from sohankshirsagar April 7, 2026 20:25

cursor bot reviewed Apr 7, 2026

View reviewed changes

drift/core/coverage_server.py Outdated Show resolved Hide resolved

ref: extract _group_arcs_by_line helper to deduplicate arc grouping

c0fac85

sohankshirsagar approved these changes Apr 7, 2026

View reviewed changes

sohil-kshirsagar merged commit 9f37908 into main Apr 8, 2026
27 checks passed

sohil-kshirsagar deleted the feat/code-coverage-tracking-poc branch April 8, 2026 00:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(coverage): Add coverage.py code coverage collection#87

feat(coverage): Add coverage.py code coverage collection#87
sohil-kshirsagar merged 24 commits intomainfrom
feat/code-coverage-tracking-poc

sohil-kshirsagar commented Apr 3, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cubic-dev-ai bot left a comment

Uh oh!

Uh oh!

Uh oh!

tusk-dev bot commented Apr 7, 2026 •

edited

Loading

Uh oh!

cursor bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

sohil-kshirsagar commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

How it works

Key design decisions

Files changed

Also includes (non-coverage)

TODOs before merge

Edge cases / gotchas

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cubic-dev-ai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

tusk-dev bot commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cursor bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sohil-kshirsagar commented Apr 3, 2026 •

edited

Loading

tusk-dev bot commented Apr 7, 2026 •

edited

Loading