Bring tracing service more in line with upstream APIs

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
Eric Peterson
2026-05-18 15:20:09 +02:00
parent d6c7805527
commit b70f13990b
15 changed files with 567 additions and 221 deletions
@@ -16,6 +16,8 @@ import { MetricsService } from '@backstage/backend-plugin-api/alpha';
import { ServiceFactory } from '@backstage/backend-plugin-api';
import { TracingService } from '@backstage/backend-plugin-api/alpha';
import { TracingServiceAttributeValue } from '@backstage/backend-plugin-api/alpha';
import { TracingServiceContextAPI } from '@backstage/backend-plugin-api/alpha';
import { TracingServicePropagationAPI } from '@backstage/backend-plugin-api/alpha';
import { TracingServiceSpan } from '@backstage/backend-plugin-api/alpha';
import { TracingServiceSpanStatus } from '@backstage/backend-plugin-api/alpha';
@@ -85,6 +87,28 @@ export class MockActionsRegistry
>(options: ActionsRegistryActionOptions<TInputSchema, TOutputSchema>): void;
}
// @alpha
export interface MockedTracingServiceContextAPI
extends TracingServiceContextAPI {
// (undocumented)
active: jest.MockedFunction<TracingServiceContextAPI['active']>;
// (undocumented)
with: jest.MockedFunction<TracingServiceContextAPI['with']>;
}
// @alpha
export interface MockedTracingServicePropagationAPI
extends TracingServicePropagationAPI {
// (undocumented)
extract: jest.MockedFunction<TracingServicePropagationAPI['extract']>;
// (undocumented)
getActiveBaggage: jest.MockedFunction<
TracingServicePropagationAPI['getActiveBaggage']
>;
// (undocumented)
getBaggage: jest.MockedFunction<TracingServicePropagationAPI['getBaggage']>;
}
// @alpha
export interface MockedTracingServiceSpan extends TracingServiceSpan {
// (undocumented)
@@ -106,17 +130,15 @@ export type ServiceMock<TService> = {
// @alpha
export interface TracingServiceMock extends TracingService {
// (undocumented)
context: MockedTracingServiceContextAPI;
// (undocumented)
factory: ServiceFactory<TracingService>;
// (undocumented)
getActiveBaggage: jest.MockedFunction<TracingService['getActiveBaggage']>;
propagation: MockedTracingServicePropagationAPI;
spans: MockedTracingServiceSpan[];
// (undocumented)
startActiveSpan: jest.MockedFunction<TracingService['startActiveSpan']>;
// (undocumented)
withPropagatedContext: jest.MockedFunction<
TracingService['withPropagatedContext']
>;
}
// @alpha (undocumented)
@@ -17,37 +17,41 @@
import { tracingServiceMock } from './TracingServiceMock';
describe('tracingServiceMock', () => {
it('parses the baggage header from withPropagatedContext and exposes it via getActiveBaggage', async () => {
it('parses the baggage header via propagation.extract and exposes it via getActiveBaggage inside context.with', async () => {
const tracing = tracingServiceMock.mock();
const seen = await tracing.withPropagatedContext(
{
baggage:
'gen_ai.conversation.id=conv-123, gen_ai.agent.id=agent-456;property=ignored',
},
() => {
const baggage = tracing.getActiveBaggage();
return {
conv: baggage?.getEntry('gen_ai.conversation.id')?.value,
agent: baggage?.getEntry('gen_ai.agent.id')?.value,
missing: baggage?.getEntry('gen_ai.missing'),
all: baggage?.getAllEntries().map(([k, v]) => [k, v.value]),
};
},
);
const ctx = tracing.propagation.extract(tracing.context.active(), {
baggage:
'gen_ai.conversation.id=conv-123, gen_ai.agent.id=agent-456;property=ignored',
});
// Baggage is reachable directly off the extracted handle.
expect(
tracing.propagation.getBaggage(ctx)?.getEntry('gen_ai.agent.id'),
).toEqual({ value: 'agent-456' });
const seen = await tracing.context.with(ctx, () => {
const baggage = tracing.propagation.getActiveBaggage();
return {
conv: baggage?.getEntry('gen_ai.conversation.id')?.value,
agent: baggage?.getEntry('gen_ai.agent.id')?.value,
missing: baggage?.getEntry('gen_ai.missing'),
all: baggage?.getAllEntries().map(([k, v]) => [k, v.value]),
};
});
expect(seen).toEqual({
conv: 'conv-123',
agent: 'agent-456',
missing: undefined,
all: [
['gen_ai.conversation.id', 'conv-123'],
['gen_ai.agent.id', 'agent-456'],
],
['gen_ai.conversation.id', { value: 'conv-123' }],
['gen_ai.agent.id', { value: 'agent-456' }],
].map(([k, v]) => [k, (v as { value: string }).value]),
});
// Baggage is scoped to the propagated callback.
expect(tracing.getActiveBaggage()).toBeUndefined();
// Baggage is scoped to the context.with callback.
expect(tracing.propagation.getActiveBaggage()).toBeUndefined();
});
it('honours mockReturnValue overrides for getActiveBaggage', async () => {
@@ -59,22 +63,25 @@ describe('tracingServiceMock', () => {
[string, { value: string }]
>,
};
tracing.getActiveBaggage.mockReturnValue(override);
tracing.propagation.getActiveBaggage.mockReturnValue(override);
expect(tracing.getActiveBaggage()).toBe(override);
await tracing.withPropagatedContext(
{ baggage: 'gen_ai.conversation.id=conv-from-header' },
() => {
// mockReturnValue takes precedence over the default header parsing.
expect(tracing.getActiveBaggage()).toBe(override);
},
);
expect(tracing.propagation.getActiveBaggage()).toBe(override);
const ctx = tracing.propagation.extract(tracing.context.active(), {
baggage: 'gen_ai.conversation.id=conv-from-header',
});
await tracing.context.with(ctx, () => {
// mockReturnValue takes precedence over the default header parsing.
expect(tracing.propagation.getActiveBaggage()).toBe(override);
});
});
it('returns undefined baggage when no baggage header is supplied', async () => {
it('returns undefined baggage when no baggage header is supplied to extract', async () => {
const tracing = tracingServiceMock.mock();
await tracing.withPropagatedContext({ traceparent: 'whatever' }, () => {
expect(tracing.getActiveBaggage()).toBeUndefined();
const ctx = tracing.propagation.extract(tracing.context.active(), {
traceparent: 'whatever',
});
await tracing.context.with(ctx, () => {
expect(tracing.propagation.getActiveBaggage()).toBeUndefined();
});
});
});
@@ -22,22 +22,38 @@ import {
TracingService,
TracingServiceAttributeValue,
TracingServiceBaggage,
TracingServiceContext,
TracingServiceContextAPI,
TracingServicePropagationAPI,
TracingServiceSpan,
TracingServiceSpanStatus,
tracingServiceRef,
} from '@backstage/backend-plugin-api/alpha';
import { tracingServiceFactory } from '@backstage/backend-defaults/alpha';
// Internal context shape used by the mock. The opaque `TracingServiceContext`
// is just this object cast to the public type.
interface MockContext {
baggage?: TracingServiceBaggage;
}
function toMockContext(ctx: TracingServiceContext): MockContext {
return ctx as unknown as MockContext;
}
function fromMockContext(ctx: MockContext): TracingServiceContext {
return ctx as unknown as TracingServiceContext;
}
// Parses the `baggage` header per the W3C Baggage member syntax,
// dropping value properties (`;property=value`). This mirrors what
// `propagation.extract` does in the real tracing service, just enough
// for tests to assert end-to-end behaviour between propagated headers
// and `getActiveBaggage()`.
function parseBaggageHeader(
headers: Record<string, string | string[] | undefined>,
carrier: Record<string, string | string[] | undefined>,
): TracingServiceBaggage | undefined {
let raw: string | undefined;
for (const [name, value] of Object.entries(headers)) {
for (const [name, value] of Object.entries(carrier)) {
if (name.toLowerCase() !== 'baggage') continue;
raw = Array.isArray(value) ? value[0] : value;
break;
@@ -72,27 +88,55 @@ export interface MockedTracingServiceSpan extends TracingServiceSpan {
setStatus: jest.Mock<void, [TracingServiceSpanStatus]>;
}
/**
* Jest-mocked counterpart of the `context` member on the
* `TracingService`.
*
* @alpha
*/
export interface MockedTracingServiceContextAPI
extends TracingServiceContextAPI {
active: jest.MockedFunction<TracingServiceContextAPI['active']>;
with: jest.MockedFunction<TracingServiceContextAPI['with']>;
}
/**
* Jest-mocked counterpart of the `propagation` member on the
* `TracingService`.
*
* @alpha
*/
export interface MockedTracingServicePropagationAPI
extends TracingServicePropagationAPI {
extract: jest.MockedFunction<TracingServicePropagationAPI['extract']>;
getBaggage: jest.MockedFunction<TracingServicePropagationAPI['getBaggage']>;
getActiveBaggage: jest.MockedFunction<
TracingServicePropagationAPI['getActiveBaggage']
>;
}
/**
* Mock for the `TracingService`. Captures every span created via
* `startActiveSpan` so tests can assert on the options passed in and the
* methods called on the span inside the callback.
*
* By default, `withPropagatedContext` parses the `baggage` header (W3C
* Baggage syntax) out of the supplied headers and makes those entries
* available via `getActiveBaggage` for the duration of the wrapped
* callback. Other propagation headers (e.g. `traceparent`) are ignored.
* Tests that need fully custom baggage can still override
* `getActiveBaggage` via `mockReturnValue` / `mockImplementation`, which
* takes precedence over the default behaviour.
* By default, `propagation.extract` parses the `baggage` header (W3C
* Baggage syntax) out of the supplied carrier and stashes the entries
* on the returned context handle. `context.with` activates that handle
* for the duration of the wrapped callback so
* `propagation.getActiveBaggage` (and `propagation.getBaggage` on the
* supplied handle) returns those entries. Other propagation fields
* (e.g. `traceparent`) are ignored. Tests that need fully custom
* baggage can still override `propagation.getActiveBaggage` via
* `mockReturnValue` / `mockImplementation`, which takes precedence over
* the default behaviour.
*
* @alpha
*/
export interface TracingServiceMock extends TracingService {
startActiveSpan: jest.MockedFunction<TracingService['startActiveSpan']>;
withPropagatedContext: jest.MockedFunction<
TracingService['withPropagatedContext']
>;
getActiveBaggage: jest.MockedFunction<TracingService['getActiveBaggage']>;
context: MockedTracingServiceContextAPI;
propagation: MockedTracingServicePropagationAPI;
/** Spans created by `startActiveSpan` calls, in order. */
spans: MockedTracingServiceSpan[];
factory: ServiceFactory<TracingService>;
@@ -113,35 +157,74 @@ export namespace tracingServiceMock {
*/
export const mock = (): TracingServiceMock => {
const spans: MockedTracingServiceSpan[] = [];
const startActiveSpan = jest.fn(async (_name, fn, _options) => {
const span: MockedTracingServiceSpan = {
setAttribute: jest.fn(),
setStatus: jest.fn(),
};
spans.push(span);
return await fn(span);
}) as TracingServiceMock['startActiveSpan'];
const startActiveSpan = jest.fn(
async (
_name: string,
optionsOrFn: unknown,
maybeFn?: (span: MockedTracingServiceSpan) => unknown,
) => {
const fn = (
typeof optionsOrFn === 'function' ? optionsOrFn : maybeFn
) as (span: MockedTracingServiceSpan) => unknown;
const span: MockedTracingServiceSpan = {
setAttribute: jest.fn(),
setStatus: jest.fn(),
};
spans.push(span);
return await fn(span);
},
) as unknown as TracingServiceMock['startActiveSpan'];
const baggageStack: Array<TracingServiceBaggage | undefined> = [];
const withPropagatedContext = jest.fn(async (headers, fn) => {
baggageStack.push(parseBaggageHeader(headers));
const contextStack: MockContext[] = [{}];
const active = jest.fn(() =>
fromMockContext(contextStack[contextStack.length - 1]),
) as MockedTracingServiceContextAPI['active'];
const withFn = jest.fn(async (ctx, fn) => {
contextStack.push(toMockContext(ctx));
try {
return await fn();
} finally {
baggageStack.pop();
contextStack.pop();
}
}) as TracingServiceMock['withPropagatedContext'];
const getActiveBaggage = jest.fn(
() => baggageStack[baggageStack.length - 1],
) as TracingServiceMock['getActiveBaggage'];
}) as MockedTracingServiceContextAPI['with'];
const service: TracingService = {
startActiveSpan,
withPropagatedContext,
const extract = jest.fn((ctx, carrier) => {
const baggage = parseBaggageHeader(carrier);
// Carry forward the parsed baggage; preserve any baggage already on the
// supplied handle if the carrier doesn't include one.
const base = toMockContext(ctx);
return fromMockContext({ baggage: baggage ?? base.baggage });
}) as MockedTracingServicePropagationAPI['extract'];
const getBaggage = jest.fn(
ctx => toMockContext(ctx).baggage,
) as MockedTracingServicePropagationAPI['getBaggage'];
const getActiveBaggage = jest.fn(
() => contextStack[contextStack.length - 1].baggage,
) as MockedTracingServicePropagationAPI['getActiveBaggage'];
const context: MockedTracingServiceContextAPI = {
active,
with: withFn,
};
const propagation: MockedTracingServicePropagationAPI = {
extract,
getBaggage,
getActiveBaggage,
};
const service: TracingService = {
startActiveSpan,
context,
propagation,
};
return Object.assign(service as TracingServiceMock, {
context,
propagation,
spans,
factory: createServiceFactory({
service: tracingServiceRef,
@@ -22,5 +22,7 @@ export {
tracingServiceMock,
type TracingServiceMock,
type MockedTracingServiceSpan,
type MockedTracingServiceContextAPI,
type MockedTracingServicePropagationAPI,
} from './TracingServiceMock';
export { type ServiceMock } from './alphaCreateServiceMock';