How the pytest Plugin Integrates Call Tracing with Test Runs in code-graph-rag

The code-graph-rag pytest plugin automatically instruments test runs to capture runtime call graphs, associating every function call with the specific test case that triggered it.

The code-graph-rag repository provides a zero-overhead tracing system that integrates directly into pytest's execution lifecycle. This article explains how the plugin in codebase_rag/trace/pytest_plugin.py wires the CallGraphTracer into your test suite, producing deterministic call-graph data tied to individual test cases for downstream analysis.

Plugin Discovery and Activation

The pytest plugin is automatically discovered via the pytest11 entry point defined in pyproject.toml. No code changes are required to enable it in an existing project.

The plugin remains inert by default until you supply the --cgr-trace command-line flag. All options are registered in pytest_addoption (lines 35-45 of pytest_plugin.py):

Flag	Purpose	Default
--cgr-trace	Enable tracing for this run	False
--cgr-trace-output	Path for the trace file	cgr-trace.jsonl (from constants.py)
--cgr-trace-repo	Repository root to limit tracing	Current working directory

Creating and Starting the Tracer

When pytest parses the command line, pytest_configure (lines 48-56) executes if tracing is enabled. This function:

1. Instantiates CallGraphTracer with the chosen repository root, filtering traces to files within that directory
2. Stores the tracer in pytest's internal stash using _STASH_KEY for retrieval across hook invocations
3. Starts the tracer by calling tracer.start(), which registers a sys.monitoring callback for PY_START events (lines 74-88 in tracer.py)

# From codebase_rag/trace/pytest_plugin.py (simplified)

def pytest_configure(config):
    if config.option.cgr_trace:
        tracer = CallGraphTracer(repo_root=config.option.cgr_trace_repo)
        config.stash[_STASH_KEY] = tracer
        tracer.start()  # Registers sys.monitoring callback

The sys.monitoring API provides low-overhead event-based tracing available in Python 3.12+, avoiding the performance penalties of traditional sys.setprofile approaches.

Associating Calls with Individual Tests

The core integration happens in pytest_runtest_protocol (lines 60-72), which wraps the entire test execution. This pytest hook executes once per test case:

Before the test runs:

• The tracer is retrieved from the stash
• tracer.set_workload(item.nodeid) is called, where item.nodeid is the test's unique identifier (e.g., test_module.py::test_function)

During execution:

• Every PY_START event fires the tracer's callback
• The callback records the caller-callee pair tagged with the current workload

After the test completes:

• The workload is cleared, ensuring isolation between tests

The set_workload mechanism (lines 63-73 in tracer.py) is critical: it enables fine-grained attribution of every function call to the specific test that triggered it. This supports downstream use cases like identifying which tests exercise particular code paths or diagnosing flaky tests through call-pattern analysis.

Parallel Execution Support with pytest-xdist

When running tests in parallel via pytest -n N, the plugin handles worker-separated trace files automatically. The _output_path helper (lines 82-89) inspects config.workerinput and appends the worker ID to the output filename:

def _output_path(config):
    path = config.option.cgr_trace_output
    if hasattr(config, "workerinput"):
        worker_id = config.workerinput["workerid"]
        # Produces: cgr-trace-0.jsonl, cgr-trace-1.jsonl, etc.

        path = path.replace(".jsonl", f"-{worker_id}.jsonl")
    return path

This design guarantees that no trace records are lost in distributed execution scenarios, with files that can be merged or processed independently.

Flushing Traces at Session Completion

The pytest_sessionfinish hook (lines 92-105) executes once per worker after all tests complete:

1. Retrieves the tracer from the stash
2. Calls tracer.stop() to deregister sys.monitoring callbacks
3. Computes the final output path via _output_path
4. Calls tracer.write(output), which serializes CallRecord objects to JSONL using the interchange format
5. Emits a terminal summary: "cgr-trace: wrote 2,384 call records to cgr-trace.jsonl"

The actual serialization logic resides in tracer.write (lines 66-77), which handles the conversion from internal data structures to the trace format defined in constants.py.

Cleanup on Unexpected Shutdown

If the pytest process terminates before sessionfinish—due to KeyboardInterrupt, unhandled exceptions, or test session failures—pytest_unconfigure (lines 108-112) ensures the tracer is properly stopped and the monitoring slot is released. This prevents resource leaks and corrupted trace states.

Running Tests with Call Tracing

Basic usage

# Enable tracing with default output

pytest -v --cgr-trace

# Custom output path

pytest --cgr-trace --cgr-trace-output=trace.jsonl

# Limit tracing to specific repository root

pytest --cgr-trace --cgr-trace-repo=/path/to/my/repo

Parallel execution

# Generates trace-0.jsonl, trace-1.jsonl, etc.

pytest -n 4 --cgr-trace

Inspecting trace output

import json

with open("cgr-trace.jsonl", "r") as f:
    for line in f:
        record = json.loads(line)
        print(f"{record['caller']['qualname']} → {record['callee']['qualname']}")
        print(f"  called {record['count']} times in workloads: {record['workloads']}")

Lifecycle Summary

Phase	Hook / Function	Action
CLI parsing	pytest_addoption	Register --cgr-trace flags
Session start	pytest_configure	Create and start CallGraphTracer
Per-test setup	pytest_runtest_protocol (pre-test)	set_workload(item.nodeid)
Per-test teardown	pytest_runtest_protocol (post-test)	Clear workload context
Worker output	_output_path	Append worker ID for xdist
Session finish	pytest_sessionfinish	Stop tracer, write trace file
Cleanup	pytest_unconfigure	Ensure profiler slot released

Key Source Files

File	Purpose
codebase_rag/trace/pytest_plugin.py	pytest lifecycle hooks and CLI integration
codebase_rag/trace/tracer.py	Core CallGraphTracer with sys.monitoring implementation
codebase_rag/constants.py	Default paths, tool name, trace format version
pyproject.toml	pytest11 entry point for automatic plugin discovery

Summary

• The pytest plugin in codebase_rag/trace/pytest_plugin.py provides seamless, opt-in call-graph tracing for any pytest test suite
• Workload tagging via set_workload(item.nodeid) enables precise attribution of function calls to specific test cases
• sys.monitoring integration delivers low-overhead tracing without code modification
• Parallel execution support via worker-specific output files ensures correctness with pytest-xdist
• The JSONL trace format supports downstream analysis for test impact analysis, coverage metrics, and flaky test diagnosis

Frequently Asked Questions

What Python version is required for the tracing to work?

The CallGraphTracer depends on sys.monitoring, which was introduced in Python 3.12. The plugin checks for availability and degrades gracefully on older versions.

How much overhead does tracing add to test runs?

According to the implementation in tracer.py, the sys.monitoring approach has minimal overhead compared to traditional profiling—typically single-digit percentage slowdowns for most workloads. The tracer only records call events, not line-by-line execution.

Can I trace only specific tests or modules?

The current implementation traces all tests when --cgr-trace is enabled. However, you can filter the resulting trace file by workloads field to isolate specific tests post-execution. Future versions may support selective tracing via pytest markers.

How do I merge traces from parallel workers?

Each worker produces a separate .jsonl file with the worker ID suffix. You can concatenate these files directly (they're line-delimited JSON) or process them with tools like jq or the code-graph-rag analysis utilities.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}