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                  — SpanRegistry: register, resolve, presets
├── 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_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:

@pytest.fixture(autouse=True)
def reset_otel_providers():
    _reset_otel_globals()
    yield
    _reset_otel_globals()
def _reset_span_group_state():
    set_enabled_span_groups(frozenset())
    SpanRegistry.clear()
    set_pp_trace_carrier(None)
@pytest.fixture(autouse=True)
def reset_span_groups():
    _reset_span_group_state()
    yield
    _reset_span_group_state()
@pytest.fixture(autouse=True)
def isolate_otel_resource_env(monkeypatch):
    monkeypatch.delenv("OTEL_RESOURCE_ATTRIBUTES", raising=False)

reset_span_groups pins an empty enabled set and then empties the SpanRegistry, in that order. Pinning first drops the stored spec, so a registration inside the next test cannot reopen groups it did not ask for — and so the clear() re-resolves nothing, which keeps an “unresolved span groups” warning out of every teardown. Reversing the two produces both failure modes: groups leaking across tests, and a spurious warning on every single test.

The SpanRegistry is process-global like the providers are, so a consumer repo writing its own conftest.py needs this reset too — without SpanRegistry.clear(), a second registration of the same namespace raises ValueError rather than leaking quietly.

OTEL_RESOURCE_ATTRIBUTES is one of the two channels build_providers resolves identity from, so an ambient value — common on an OTel-instrumented CI runner — would otherwise supply nv.dl.rank or service.instance.id to tests written to assert the no-rank path.

NeMo Lens ships no groups of its own, so a test that needs one must register it. The demo_groups fixture in conftest.py registers a small set for exactly this:

def test_something(demo_groups):
    cfg = NemoLensConfig(enabled=True, span_groups="per_step", exporter="console")
    setup_telemetry(cfg)
    assert is_span_group_enabled("step")

For a test that wants a specific set with no registration at all, set_enabled_span_groups(...) — a top-level export from nemo.lens — pins one directly.

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

def test_my_instrumentation():
    from tests.conftest import InMemorySpanExporter
    exporter = InMemorySpanExporter()
    cfg = NemoLensConfig(enabled=True, exporter="console")
    setup_telemetry(cfg, span_exporter=exporter)
    with managed_span('job', 'my.op') as span:
        ...
    spans = exporter.get_finished_spans()
    assert len(spans) == 1
    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.

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

Test Double-Init

def test_double_init_raises():
    cfg = NemoLensConfig(enabled=True, exporter="console")
    setup_telemetry(cfg)
    with pytest.raises(RuntimeError, match="already been initialized"):
        setup_telemetry(cfg)

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

def test_every_process_exports():
    cfg = NemoLensConfig(enabled=True, exporter="console")
    for _ in range(4):
        handle = setup_telemetry(cfg, _allow_reinit=True)
        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.

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.