Test

The NeMo Lens test suite is small and focused. Over 180 tests cover every public API, with a strong emphasis on:

• State isolation ensures that tests touching global OTel state reset it before and after running.
• No-op equivalence ensures that the behavior of fallbacks.py matches the real API.
• Configuration edge cases cover every environment variable combination and validation path.

Run Tests

$
# Full suite
$
pytest
$
$
# Single file
$
pytest tests/test_helpers.py
$
$
# Single test
$
pytest tests/test_helpers.py::TestManagedSpan::test_enabled_group_creates_span
$
$
# With coverage
$
pytest --cov=nemo.lens --cov-report=term-missing

Layout

tests/
├── conftest.py                     — shared fixtures (state reset, InMemorySpanExporter)
├── test_config.py                  — NemoLensConfig, from_env, validation
├── test_state.py                   — is_span_group_enabled, thread safety
├── test_groups.py                  — SpanGroup.resolve, preset handling
├── test_helpers.py                 — managed_span, trace_fn, span_cm, attribute safety
├── test_handle.py                  — setup_telemetry, TelemetryHandle, double-init guard
├── test_providers.py               — build_providers, custom exporters
├── test_sampling_integration.py    — RankAwareSampler, integration with TracerProvider
├── test_strategies.py              — register/unregister/registered export strategies
├── test_distributed.py             — broadcast_trace_context, create_linked_span
├── test_propagation.py             — inject_context, extract_context
├── test_resources.py               — SLURM, K8s, local detection
├── test_instruments.py             — metric recording functions
├── test_fallbacks.py               — fallback no-op correctness
└── test_e2e.py                     — end-to-end with real SDK + InMemorySpanExporter

Global State Isolation

The OTel SDK stores providers globally (trace._TRACER_PROVIDER, metrics._METER_PROVIDER). NeMo Lens stores enabled span groups globally. Tests that touch these must reset between runs; otherwise, the second test inherits the state of the first test.

conftest.py has three autouse fixtures:

1
@pytest.fixture(autouse=True)
2
def reset_otel_providers():
3
    _reset_otel_globals()
4
    yield
5
    _reset_otel_globals()
6
7
@pytest.fixture(autouse=True)
8
def reset_span_groups():
9
    set_enabled_span_groups(frozenset())
10
    set_pp_trace_carrier(None)
11
    yield
12
    set_enabled_span_groups(frozenset())
13
    set_pp_trace_carrier(None)

Because reset_span_groups clears the enabled set before every test, a test that needs a group active must opt in explicitly. The set_enabled_span_groups(...) function, which is a top-level export from nemo.lens, provides the escape hatch for enabling groups inside a test.

A third autouse fixture, reset_strategy_registry, snapshots nemo.lens.strategies._REGISTRY (under _REGISTRY_LOCK) before each test and restores it afterward, so custom export strategies registered through register_export_strategy do not leak between tests:

1
@pytest.fixture(autouse=True)
2
def reset_strategy_registry():
3
    from nemo.lens.strategies import _REGISTRY, _REGISTRY_LOCK
4
    with _REGISTRY_LOCK:
5
        snapshot = dict(_REGISTRY)
6
    yield
7
    with _REGISTRY_LOCK:
8
        _REGISTRY.clear()
9
        _REGISTRY.update(snapshot)

_reset_otel_globals() resets five pieces of state:

1
_trace_mod._TRACER_PROVIDER = None
2
_trace_mod._TRACER_PROVIDER_SET_ONCE = Once()
3
_metrics_mod._METER_PROVIDER = None
4
_metrics_mod._METER_PROVIDER_SET_ONCE = Once()
5
_handle_mod._INITIALIZED = False     # lens's own double-init flag

The Once() pointers are the OTel SDK’s internal “was this set?” flag. Without resetting them, setup_telemetry in a test would install a new provider, but the SDK would log “provider already set” and silently use the previous one.

Capture Spans

Tests that assert on span content use InMemorySpanExporter (shipped in conftest.py):

1
def test_my_instrumentation():
2
    from tests.conftest import InMemorySpanExporter
3
    exporter = InMemorySpanExporter()
4
5
    cfg = NemoLensConfig(enabled=True, exporter="console")
6
    setup_telemetry(cfg, rank=0, world_size=1, span_exporter=exporter)
7
8
    with managed_span('job', 'my.op') as span:
9
        ...
10
11
    spans = exporter.get_finished_spans()
12
    assert len(spans) == 1
13
    assert spans[0].name == 'my.op'

For metrics, use InMemoryMetricReader from the OTel SDK.

Test Fallbacks

tests/test_fallbacks.py asserts that nemo.lens.fallbacks signatures match the real API and behave as no-ops. Whenever you add a parameter to managed_span or trace_fn, also add it to fallbacks.py and extend the test.

1
def test_managed_span_accepts_kwargs():
2
    # Real API accepts arbitrary kwargs — fallback must too
3
    with managed_span("group", "name", iteration=1, loss=0.5) as span:
4
        assert span is None

Test Double-Init

1
def test_double_init_raises():
2
    cfg = NemoLensConfig(enabled=True, exporter="console")
3
    setup_telemetry(cfg, rank=0, world_size=1)
4
    with pytest.raises(RuntimeError, match="already been initialized"):
5
        setup_telemetry(cfg, rank=0, world_size=1)

Tests that legitimately need to call setup_telemetry multiple times in one test (e.g., simulating multiple ranks) pass _allow_reinit=True:

1
def test_all_ranks_export():
2
    cfg = NemoLensConfig(enabled=True, export_strategy="all_ranks", exporter="console")
3
    for rank in range(4):
4
        handle = setup_telemetry(cfg, rank=rank, world_size=4, _allow_reinit=True)
5
        assert handle.is_exporting

Test Distributed Helpers

broadcast_trace_context uses torch.distributed, which cannot run in a single-process test without mocks. The distributed tests use torch.distributed.init_process_group(backend='gloo', ...) with a single rank; therefore, the broadcast becomes a no-op, but the code path exercises correctly.

For genuinely multi-rank behavior, tests would need to spawn subprocesses; currently, the single-rank path and manual carrier construction cover the contract.

Test the OTel Interface of RankAwareSampler

1
def test_sampler_returns_proper_sampling_result():
2
    from opentelemetry.sdk.trace.sampling import Decision
3
    sampler = RankAwareSampler(rank=0, world_size=4, sample_rate=1.0)
4
    result = sampler.should_sample(parent_context=None, trace_id=12345, name="test")
5
    assert result.decision == Decision.RECORD_AND_SAMPLE

The sampler is wrapped in a try/except block inside should_sample to fall back to bool if the SDK is not installed, which is covered by a separate test.

Lint Code

$
ruff check src tests --fix
$
ruff format src tests

Pre-commit runs both (the ruff and ruff-format hooks in .pre-commit-config.yaml). CI runs pre-commit run --all-files and rejects PRs that fail it.

What Is Not Tested

• Actual export to a collector. That is the SDK’s job; mocking it correctly requires more effort than it is worth.
• Long-running performance. The tests/ directory exercises correctness, not throughput.
• Integration with consumer libraries. These libraries have their own test suites (Megatron-LM/tests/unit_tests/telemetry/, etc.).

When adding features that interact with a consumer, add a corresponding test in the consumer repository. NeMo Lens tests must remain self-contained.