Bring tracing service more in line with upstream APIs
Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
@@ -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';
|
||||
|
||||
Reference in New Issue
Block a user