Skip to content

Commit 4b39078

Browse files
chargomechargome
committed
useFlushIfServerless
1 parent b083516 commit 4b39078

File tree

2 files changed

+149
-59
lines changed

2 files changed

+149
-59
lines changed

‎packages/react-router/src/server/wrapSentryHandleRequest.ts

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,12 @@ import { context } from '@opentelemetry/api';
22
import { getRPCMetadata, RPCType } from '@opentelemetry/core';
33
import { ATTR_HTTP_ROUTE } from '@opentelemetry/semantic-conventions';
44
import {
5-
flush,
5+
flushIfServerless,
66
getActiveSpan,
77
getRootSpan,
88
getTraceMetaTags,
99
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
1010
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
11-
vercelWaitUntil,
1211
} from '@sentry/core';
1312
import type { AppLoadContext, EntryContext } from 'react-router';
1413
import type { PassThrough } from 'stream';
@@ -63,7 +62,7 @@ export function wrapSentryHandleRequest(originalHandle: OriginalHandleRequest):
6362
try {
6463
return await originalHandle(request, responseStatusCode, responseHeaders, routerContext, loadContext);
6564
} finally {
66-
vercelWaitUntil(flush());
65+
await flushIfServerless();
6766
}
6867
};
6968
}

‎packages/react-router/test/server/wrapSentryHandleRequest.test.ts

Lines changed: 147 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
11
import { RPCType } from '@opentelemetry/core';
22
import { ATTR_HTTP_ROUTE } from '@opentelemetry/semantic-conventions';
33
import {
4-
flush,
4+
flushIfServerless,
55
getActiveSpan,
66
getRootSpan,
77
getTraceMetaTags,
88
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
99
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
10-
vercelWaitUntil,
1110
} from '@sentry/core';
1211
import { PassThrough } from 'stream';
1312
import { beforeEach, describe, expect, test, vi } from 'vitest';
@@ -17,73 +16,20 @@ vi.mock('@opentelemetry/core', () => ({
1716
RPCType: { HTTP: 'http' },
1817
getRPCMetadata: vi.fn(),
1918
}));
20-
2119
vi.mock('@sentry/core', () => ({
2220
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE: 'sentry.source',
2321
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN: 'sentry.origin',
2422
getActiveSpan: vi.fn(),
2523
getRootSpan: vi.fn(),
2624
getTraceMetaTags: vi.fn(),
27-
flush: vi.fn(),
28-
vercelWaitUntil: vi.fn(),
25+
flushIfServerless: vi.fn(),
2926
}));
3027

3128
describe('wrapSentryHandleRequest', () => {
3229
beforeEach(() => {
3330
vi.clearAllMocks();
3431
});
3532

36-
test('should call flush on successful execution', async () => {
37-
const originalHandler = vi.fn().mockResolvedValue('success response');
38-
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
39-
40-
const request = new Request('https://example.com');
41-
const responseStatusCode = 200;
42-
const responseHeaders = new Headers();
43-
const routerContext = { staticHandlerContext: { matches: [] } } as any;
44-
const loadContext = {} as any;
45-
46-
await wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext);
47-
48-
expect(vercelWaitUntil).toHaveBeenCalledWith(flush());
49-
expect(flush).toHaveBeenCalled();
50-
});
51-
52-
test('should call flush even when original handler throws an error', async () => {
53-
const mockError = new Error('Handler failed');
54-
const originalHandler = vi.fn().mockRejectedValue(mockError);
55-
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
56-
57-
const request = new Request('https://example.com');
58-
const responseStatusCode = 200;
59-
const responseHeaders = new Headers();
60-
const routerContext = { staticHandlerContext: { matches: [] } } as any;
61-
const loadContext = {} as any;
62-
63-
await expect(
64-
wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext),
65-
).rejects.toThrow('Handler failed');
66-
67-
expect(vercelWaitUntil).toHaveBeenCalledWith(flush());
68-
expect(flush).toHaveBeenCalled();
69-
});
70-
71-
test('should propagate errors from original handler', async () => {
72-
const mockError = new Error('Test error');
73-
const originalHandler = vi.fn().mockRejectedValue(mockError);
74-
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
75-
76-
const request = new Request('https://example.com');
77-
const responseStatusCode = 500;
78-
const responseHeaders = new Headers();
79-
const routerContext = { staticHandlerContext: { matches: [] } } as any;
80-
const loadContext = {} as any;
81-
82-
await expect(wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext)).rejects.toBe(
83-
mockError,
84-
);
85-
});
86-
8733
test('should call original handler with same parameters', async () => {
8834
const originalHandler = vi.fn().mockResolvedValue('original response');
8935
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
@@ -179,6 +125,151 @@ describe('wrapSentryHandleRequest', () => {
179125

180126
expect(getRPCMetadata).not.toHaveBeenCalled();
181127
});
128+
129+
test('should call flushIfServerless on successful execution', async () => {
130+
const originalHandler = vi.fn().mockResolvedValue('success response');
131+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
132+
133+
const request = new Request('https://example.com');
134+
const responseStatusCode = 200;
135+
const responseHeaders = new Headers();
136+
const routerContext = { staticHandlerContext: { matches: [] } } as any;
137+
const loadContext = {} as any;
138+
139+
await wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext);
140+
141+
expect(flushIfServerless).toHaveBeenCalled();
142+
});
143+
144+
test('should call flushIfServerless even when original handler throws an error', async () => {
145+
const mockError = new Error('Handler failed');
146+
const originalHandler = vi.fn().mockRejectedValue(mockError);
147+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
148+
149+
const request = new Request('https://example.com');
150+
const responseStatusCode = 200;
151+
const responseHeaders = new Headers();
152+
const routerContext = { staticHandlerContext: { matches: [] } } as any;
153+
const loadContext = {} as any;
154+
155+
await expect(
156+
wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext),
157+
).rejects.toThrow('Handler failed');
158+
159+
expect(flushIfServerless).toHaveBeenCalled();
160+
});
161+
162+
test('should propagate errors from original handler', async () => {
163+
const mockError = new Error('Test error');
164+
const originalHandler = vi.fn().mockRejectedValue(mockError);
165+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
166+
167+
const request = new Request('https://example.com');
168+
const responseStatusCode = 500;
169+
const responseHeaders = new Headers();
170+
const routerContext = { staticHandlerContext: { matches: [] } } as any;
171+
const loadContext = {} as any;
172+
173+
await expect(wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext)).rejects.toBe(
174+
mockError,
175+
);
176+
});
177+
});
178+
179+
test('should call original handler with same parameters', async () => {
180+
const originalHandler = vi.fn().mockResolvedValue('original response');
181+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
182+
183+
const request = new Request('https://taco.burrito');
184+
const responseStatusCode = 200;
185+
const responseHeaders = new Headers();
186+
const routerContext = { staticHandlerContext: { matches: [] } } as any;
187+
const loadContext = {} as any;
188+
189+
const result = await wrappedHandler(request, responseStatusCode, responseHeaders, routerContext, loadContext);
190+
191+
expect(originalHandler).toHaveBeenCalledWith(
192+
request,
193+
responseStatusCode,
194+
responseHeaders,
195+
routerContext,
196+
loadContext,
197+
);
198+
expect(result).toBe('original response');
199+
});
200+
201+
test('should set span attributes when parameterized path exists and active span exists', async () => {
202+
const originalHandler = vi.fn().mockResolvedValue('test');
203+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
204+
205+
const mockActiveSpan = {};
206+
const mockRootSpan = { setAttributes: vi.fn() };
207+
const mockRpcMetadata = { type: RPCType.HTTP, route: '/some-path' };
208+
209+
(getActiveSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockActiveSpan);
210+
(getRootSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockRootSpan);
211+
const getRPCMetadata = vi.fn().mockReturnValue(mockRpcMetadata);
212+
(vi.importActual('@opentelemetry/core') as unknown as { getRPCMetadata: typeof getRPCMetadata }).getRPCMetadata =
213+
getRPCMetadata;
214+
215+
const routerContext = {
216+
staticHandlerContext: {
217+
matches: [{ route: { path: 'some-path' } }],
218+
},
219+
} as any;
220+
221+
await wrappedHandler(new Request('https://nacho.queso'), 200, new Headers(), routerContext, {} as any);
222+
223+
expect(mockRootSpan.setAttributes).toHaveBeenCalledWith({
224+
[ATTR_HTTP_ROUTE]: '/some-path',
225+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
226+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.http.react-router.request-handler',
227+
});
228+
expect(mockRpcMetadata.route).toBe('/some-path');
229+
});
230+
231+
test('should not set span attributes when parameterized path does not exist', async () => {
232+
const mockActiveSpan = {};
233+
const mockRootSpan = { setAttributes: vi.fn() };
234+
235+
(getActiveSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockActiveSpan);
236+
(getRootSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(mockRootSpan);
237+
238+
const originalHandler = vi.fn().mockResolvedValue('test');
239+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
240+
241+
const routerContext = {
242+
staticHandlerContext: {
243+
matches: [],
244+
},
245+
} as any;
246+
247+
await wrappedHandler(new Request('https://guapo.chulo'), 200, new Headers(), routerContext, {} as any);
248+
249+
expect(mockRootSpan.setAttributes).not.toHaveBeenCalled();
250+
});
251+
252+
test('should not set span attributes when active span does not exist', async () => {
253+
const originalHandler = vi.fn().mockResolvedValue('test');
254+
const wrappedHandler = wrapSentryHandleRequest(originalHandler);
255+
256+
const mockRpcMetadata = { type: RPCType.HTTP, route: '/some-path' };
257+
258+
(getActiveSpan as unknown as ReturnType<typeof vi.fn>).mockReturnValue(null);
259+
260+
const getRPCMetadata = vi.fn().mockReturnValue(mockRpcMetadata);
261+
(vi.importActual('@opentelemetry/core') as unknown as { getRPCMetadata: typeof getRPCMetadata }).getRPCMetadata =
262+
getRPCMetadata;
263+
264+
const routerContext = {
265+
staticHandlerContext: {
266+
matches: [{ route: { path: 'some-path' } }],
267+
},
268+
} as any;
269+
270+
await wrappedHandler(new Request('https://tio.pepe'), 200, new Headers(), routerContext, {} as any);
271+
272+
expect(getRPCMetadata).not.toHaveBeenCalled();
182273
});
183274

184275
describe('getMetaTagTransformer', () => {

0 commit comments

Comments
 (0)