From 0eb536eb0d399e409ff9bdf81d4b2333bef0ca83 Mon Sep 17 00:00:00 2001 From: JPeer264 Date: Wed, 18 Feb 2026 12:33:58 -0100 Subject: [PATCH] feat(cloudflare): Instrument async KV API --- .../cloudflare-workers/src/index.ts | 15 +- .../cloudflare-workers/tests/index.test.ts | 19 +- .../instrumentDurableObjectStorage.ts | 50 +++ .../cloudflare/src/utils/instrumentContext.ts | 29 +- .../cloudflare/src/wrapMethodWithSentry.ts | 10 +- .../cloudflare/test/instrumentContext.test.ts | 87 +++++ .../instrumentDurableObjectStorage.test.ts | 212 ++++++++++++ .../test/wrapMethodWithSentry.test.ts | 308 ++++++++++++++++++ 8 files changed, 723 insertions(+), 7 deletions(-) create mode 100644 packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts create mode 100644 packages/cloudflare/test/instrumentDurableObjectStorage.test.ts create mode 100644 packages/cloudflare/test/wrapMethodWithSentry.test.ts diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts b/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts index ab438432a004..cc71748c44f8 100644 --- a/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts +++ b/dev-packages/e2e-tests/test-applications/cloudflare-workers/src/index.ts @@ -20,17 +20,26 @@ class MyDurableObjectBase extends DurableObject { } async fetch(request: Request) { - const { pathname } = new URL(request.url); - switch (pathname) { + const url = new URL(request.url); + switch (url.pathname) { case '/throwException': { await this.throwException(); break; } - case '/ws': + case '/ws': { const webSocketPair = new WebSocketPair(); const [client, server] = Object.values(webSocketPair); this.ctx.acceptWebSocket(server); return new Response(null, { status: 101, webSocket: client }); + } + case '/storage/put': { + await this.ctx.storage.put('test-key', 'test-value'); + return new Response('Stored'); + } + case '/storage/get': { + const value = await this.ctx.storage.get('test-key'); + return new Response(`Got: ${value}`); + } } return new Response('DO is fine'); } diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts b/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts index 8c09693c81ed..4235ca7d17cc 100644 --- a/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts +++ b/dev-packages/e2e-tests/test-applications/cloudflare-workers/tests/index.test.ts @@ -1,5 +1,5 @@ import { expect, test } from '@playwright/test'; -import { waitForError, waitForRequest } from '@sentry-internal/test-utils'; +import { waitForError, waitForRequest, waitForTransaction } from '@sentry-internal/test-utils'; import { SDK_VERSION } from '@sentry/cloudflare'; import { WebSocket } from 'ws'; @@ -82,3 +82,20 @@ test('sends user-agent header with SDK name and version in envelope requests', a 'user-agent': `sentry.javascript.cloudflare/${SDK_VERSION}`, }); }); + +test('Storage operations create spans in Durable Object transactions', async ({ baseURL }) => { + const transactionWaiter = waitForTransaction('cloudflare-workers', event => { + return event.spans?.some(span => span.op === 'db' && span.description === 'durable_object_storage_put') ?? false; + }); + + const response = await fetch(`${baseURL}/pass-to-object/storage/put`); + expect(response.status).toBe(200); + + const transaction = await transactionWaiter; + const putSpan = transaction.spans?.find(span => span.description === 'durable_object_storage_put'); + + expect(putSpan).toBeDefined(); + expect(putSpan?.op).toBe('db'); + expect(putSpan?.data?.['db.system.name']).toBe('cloudflare.durable_object.storage'); + expect(putSpan?.data?.['db.operation.name']).toBe('put'); +}); diff --git a/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts new file mode 100644 index 000000000000..29d47eb481f3 --- /dev/null +++ b/packages/cloudflare/src/instrumentations/instrumentDurableObjectStorage.ts @@ -0,0 +1,50 @@ +import type { DurableObjectStorage } from '@cloudflare/workers-types'; +import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startSpan } from '@sentry/core'; + +const STORAGE_METHODS_TO_INSTRUMENT = ['get', 'put', 'delete', 'list'] as const; + +type StorageMethod = (typeof STORAGE_METHODS_TO_INSTRUMENT)[number]; + +/** + * Instruments DurableObjectStorage methods with Sentry spans. + * + * Wraps the following async methods: + * - get, put, delete, list (KV API) + * + * @param storage - The DurableObjectStorage instance to instrument + * @returns An instrumented DurableObjectStorage instance + */ +export function instrumentDurableObjectStorage(storage: DurableObjectStorage): DurableObjectStorage { + return new Proxy(storage, { + get(target, prop, receiver) { + const original = Reflect.get(target, prop, receiver); + + if (typeof original !== 'function') { + return original; + } + + const methodName = prop as string; + if (!STORAGE_METHODS_TO_INSTRUMENT.includes(methodName as StorageMethod)) { + return (original as (...args: unknown[]) => unknown).bind(target); + } + + return function (this: unknown, ...args: unknown[]) { + return startSpan( + { + // Use underscore naming to match Cloudflare's native instrumentation (e.g., "durable_object_storage_get") + name: `durable_object_storage_${methodName}`, + op: 'db', + attributes: { + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.db.cloudflare.durable_object', + 'db.system.name': 'cloudflare.durable_object.storage', + 'db.operation.name': methodName, + }, + }, + () => { + return (original as (...args: unknown[]) => unknown).apply(target, args); + }, + ); + }; + }, + }); +} diff --git a/packages/cloudflare/src/utils/instrumentContext.ts b/packages/cloudflare/src/utils/instrumentContext.ts index 5f5d0577f64e..a8c04c318a2d 100644 --- a/packages/cloudflare/src/utils/instrumentContext.ts +++ b/packages/cloudflare/src/utils/instrumentContext.ts @@ -1,4 +1,5 @@ -import { type DurableObjectState, type ExecutionContext } from '@cloudflare/workers-types'; +import { type DurableObjectState, type DurableObjectStorage, type ExecutionContext } from '@cloudflare/workers-types'; +import { instrumentDurableObjectStorage } from '../instrumentations/instrumentDurableObjectStorage'; type ContextType = ExecutionContext | DurableObjectState; type OverridesStore = Map unknown>; @@ -8,6 +9,8 @@ type OverridesStore = Map(ctx: T): T { {} as PropertyDescriptorMap, ); + // Check if this is a DurableObjectState context with a storage property + // If so, wrap the storage with instrumentation + if ('storage' in ctx && ctx.storage) { + const originalStorage = ctx.storage; + let instrumentedStorage: DurableObjectStorage | undefined; + descriptors.storage = { + configurable: true, + enumerable: true, + get: () => { + if (!instrumentedStorage) { + instrumentedStorage = instrumentDurableObjectStorage(originalStorage); + } + return instrumentedStorage; + }, + }; + // Expose the original uninstrumented storage for internal Sentry operations + // This avoids creating spans for internal storage operations + descriptors.originalStorage = { + configurable: true, + enumerable: false, + get: () => originalStorage, + }; + } + return Object.create(ctx, descriptors); } diff --git a/packages/cloudflare/src/wrapMethodWithSentry.ts b/packages/cloudflare/src/wrapMethodWithSentry.ts index e3a4b1ce2b0a..f0fe3c83f5e0 100644 --- a/packages/cloudflare/src/wrapMethodWithSentry.ts +++ b/packages/cloudflare/src/wrapMethodWithSentry.ts @@ -1,3 +1,4 @@ +import type { DurableObjectStorage } from '@cloudflare/workers-types'; import { captureException, flush, @@ -14,6 +15,11 @@ import type { CloudflareOptions } from './client'; import { isInstrumented, markAsInstrumented } from './instrument'; import { init } from './sdk'; +/** Extended DurableObjectState with originalStorage exposed by instrumentContext */ +interface InstrumentedDurableObjectState extends DurableObjectState { + originalStorage?: DurableObjectStorage; +} + type MethodWrapperOptions = { spanName?: string; spanOp?: string; @@ -58,13 +64,13 @@ export function wrapMethodWithSentry( // In certain situations, the passed context can become undefined. // For example, for Astro while prerendering pages at build time. // see: https://github.com/getsentry/sentry-javascript/issues/13217 - const context = wrapperOptions.context as ExecutionContext | undefined; + const context = wrapperOptions.context as InstrumentedDurableObjectState | undefined; const waitUntil = context?.waitUntil?.bind?.(context); const currentClient = scope.getClient(); if (!currentClient) { - const client = init({ ...wrapperOptions.options, ctx: context }); + const client = init({ ...wrapperOptions.options, ctx: context as unknown as ExecutionContext | undefined }); scope.setClient(client); } diff --git a/packages/cloudflare/test/instrumentContext.test.ts b/packages/cloudflare/test/instrumentContext.test.ts index 6cca64fd4bb1..4304a303a155 100644 --- a/packages/cloudflare/test/instrumentContext.test.ts +++ b/packages/cloudflare/test/instrumentContext.test.ts @@ -46,6 +46,57 @@ describe('instrumentContext', () => { const instrumented = instrumentContext(context); expect(instrumented[s]).toBe(context[s]); }); + + describe('DurableObjectState storage instrumentation', () => { + it('instruments storage property', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // The storage property should be instrumented (wrapped) + expect(instrumented.storage).toBeDefined(); + // The instrumented storage should not be the same reference + expect(instrumented.storage).not.toBe(mockStorage); + }); + + it('exposes originalStorage as the uninstrumented storage', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // originalStorage should be the original uninstrumented storage + expect(instrumented.originalStorage).toBe(mockStorage); + }); + + it('originalStorage is not enumerable', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // originalStorage should not appear in Object.keys + expect(Object.keys(instrumented)).not.toContain('originalStorage'); + }); + + it('returns instrumented storage lazily', () => { + const mockStorage = createMockStorage(); + const context = makeDurableObjectStateMock(mockStorage); + const instrumented = instrumentContext(context); + + // Access storage twice to ensure memoization + const storage1 = instrumented.storage; + const storage2 = instrumented.storage; + + expect(storage1).toBe(storage2); + }); + + it('handles context without storage property', () => { + const context = makeExecutionContextMock(); + const instrumented = instrumentContext(context) as any; + + // Should not have originalStorage if no storage property + expect(instrumented.originalStorage).toBeUndefined(); + }); + }); }); function makeExecutionContextMock() { @@ -54,3 +105,39 @@ function makeExecutionContextMock() { passThroughOnException: vi.fn(), } as unknown as Mocked; } + +function makeDurableObjectStateMock(storage?: any) { + return { + waitUntil: vi.fn(), + blockConcurrencyWhile: vi.fn(), + id: { toString: () => 'test-id', equals: vi.fn(), name: 'test' }, + storage: storage || createMockStorage(), + acceptWebSocket: vi.fn(), + getWebSockets: vi.fn().mockReturnValue([]), + setWebSocketAutoResponse: vi.fn(), + getWebSocketAutoResponse: vi.fn(), + getWebSocketAutoResponseTimestamp: vi.fn(), + setHibernatableWebSocketEventTimeout: vi.fn(), + getHibernatableWebSocketEventTimeout: vi.fn(), + getTags: vi.fn().mockReturnValue([]), + abort: vi.fn(), + } as any; +} + +function createMockStorage(): any { + return { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + list: vi.fn().mockResolvedValue(new Map()), + getAlarm: vi.fn().mockResolvedValue(null), + setAlarm: vi.fn().mockResolvedValue(undefined), + deleteAlarm: vi.fn().mockResolvedValue(undefined), + deleteAll: vi.fn().mockResolvedValue(undefined), + sync: vi.fn().mockResolvedValue(undefined), + transaction: vi.fn().mockImplementation(async (cb: () => unknown) => cb()), + sql: { + exec: vi.fn(), + }, + }; +} diff --git a/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts new file mode 100644 index 000000000000..11c3228f905b --- /dev/null +++ b/packages/cloudflare/test/instrumentDurableObjectStorage.test.ts @@ -0,0 +1,212 @@ +import * as sentryCore from '@sentry/core'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { instrumentDurableObjectStorage } from '../src/instrumentations/instrumentDurableObjectStorage'; + +vi.mock('@sentry/core', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + startSpan: vi.fn((opts, callback) => callback()), + getActiveSpan: vi.fn(), + }; +}); + +describe('instrumentDurableObjectStorage', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('get', () => { + it('instruments get with single key', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.get('myKey'); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_get', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'get', + }), + }, + expect.any(Function), + ); + }); + + it('instruments get with array of keys', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.get(['key1', 'key2', 'key3']); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_get', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'get', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('put', () => { + it('instruments put with single key', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.put('myKey', 'myValue'); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_put', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'put', + }), + }, + expect.any(Function), + ); + }); + + it('instruments put with object entries', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.put({ key1: 'val1', key2: 'val2' }); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_put', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'put', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('delete', () => { + it('instruments delete with single key', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.delete('myKey'); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_delete', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'delete', + }), + }, + expect.any(Function), + ); + }); + + it('instruments delete with array of keys', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.delete(['key1', 'key2']); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_delete', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'delete', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('list', () => { + it('instruments list', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.list(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + { + name: 'durable_object_storage_list', + op: 'db', + attributes: expect.objectContaining({ + 'db.operation.name': 'list', + }), + }, + expect.any(Function), + ); + }); + }); + + describe('non-instrumented methods', () => { + it('does not instrument alarm methods', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.getAlarm(); + await instrumented.setAlarm(Date.now() + 1000); + await instrumented.deleteAlarm(); + + expect(sentryCore.startSpan).not.toHaveBeenCalled(); + }); + + it('does not instrument deleteAll, sync, transaction', async () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await instrumented.deleteAll(); + await instrumented.sync(); + await instrumented.transaction(async txn => txn); + + expect(sentryCore.startSpan).not.toHaveBeenCalled(); + }); + + it('does not instrument sql property', () => { + const mockStorage = createMockStorage(); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + // sql is a property, not a method we instrument + expect(instrumented.sql).toBe(mockStorage.sql); + }); + }); + + describe('error handling', () => { + it('propagates errors from storage operations', async () => { + const mockStorage = createMockStorage(); + mockStorage.get = vi.fn().mockRejectedValue(new Error('Storage error')); + const instrumented = instrumentDurableObjectStorage(mockStorage); + + await expect(instrumented.get('myKey')).rejects.toThrow('Storage error'); + }); + }); +}); + +function createMockStorage(): any { + return { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + list: vi.fn().mockResolvedValue(new Map()), + getAlarm: vi.fn().mockResolvedValue(null), + setAlarm: vi.fn().mockResolvedValue(undefined), + deleteAlarm: vi.fn().mockResolvedValue(undefined), + deleteAll: vi.fn().mockResolvedValue(undefined), + sync: vi.fn().mockResolvedValue(undefined), + transaction: vi.fn().mockImplementation(async (cb: () => unknown) => cb()), + sql: { + exec: vi.fn(), + }, + }; +} diff --git a/packages/cloudflare/test/wrapMethodWithSentry.test.ts b/packages/cloudflare/test/wrapMethodWithSentry.test.ts new file mode 100644 index 000000000000..3acafaba9b33 --- /dev/null +++ b/packages/cloudflare/test/wrapMethodWithSentry.test.ts @@ -0,0 +1,308 @@ +import * as sentryCore from '@sentry/core'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { isInstrumented } from '../src/instrument'; +import { wrapMethodWithSentry } from '../src/wrapMethodWithSentry'; + +// Mock the SDK init to avoid actual SDK initialization +vi.mock('../src/sdk', () => ({ + init: vi.fn(() => ({ + getOptions: () => ({}), + on: vi.fn(), + })), +})); + +// Mock sentry/core functions +vi.mock('@sentry/core', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + getClient: vi.fn(), + withIsolationScope: vi.fn((callback: (scope: any) => any) => callback(createMockScope())), + withScope: vi.fn((callback: (scope: any) => any) => callback(createMockScope())), + startSpan: vi.fn((opts, callback) => callback(createMockSpan())), + captureException: vi.fn(), + flush: vi.fn().mockResolvedValue(true), + getActiveSpan: vi.fn(), + }; +}); + +function createMockScope() { + return { + getClient: vi.fn(), + setClient: vi.fn(), + }; +} + +function createMockSpan() { + return { + setAttribute: vi.fn(), + setAttributes: vi.fn(), + spanContext: vi.fn().mockReturnValue({ + traceId: 'test-trace-id-12345678901234567890', + spanId: 'test-span-id', + }), + }; +} + +function createMockContext(options: { hasStorage?: boolean; hasWaitUntil?: boolean } = {}) { + const mockStorage = { + get: vi.fn().mockResolvedValue(undefined), + put: vi.fn().mockResolvedValue(undefined), + delete: vi.fn().mockResolvedValue(false), + }; + + return { + waitUntil: options.hasWaitUntil !== false ? vi.fn() : undefined, + storage: options.hasStorage !== false ? mockStorage : undefined, + originalStorage: options.hasStorage !== false ? mockStorage : undefined, + } as any; +} + +describe('wrapMethodWithSentry', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('basic wrapping', () => { + it('wraps a sync method and returns its result', () => { + const handler = vi.fn().mockReturnValue('sync-result'); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + wrapped(); + + expect(handler).toHaveBeenCalled(); + }); + + it('wraps an async method and returns a promise', async () => { + const handler = vi.fn().mockResolvedValue('async-result'); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(handler).toHaveBeenCalled(); + }); + + it('marks handler as instrumented', () => { + const handler = vi.fn(); + const options = { + options: {}, + context: createMockContext(), + }; + + expect(isInstrumented(handler)).toBeUndefined(); + + wrapMethodWithSentry(options, handler); + + expect(isInstrumented(handler)).toBe(true); + }); + + it('does not re-wrap already instrumented handler', () => { + const handler = vi.fn(); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped1 = wrapMethodWithSentry(options, handler); + const wrapped2 = wrapMethodWithSentry(options, wrapped1); + + // Should return the same wrapped function + expect(wrapped2).toBe(wrapped1); + }); + + it('does not mark handler when noMark is true', () => { + const handler = vi.fn(); + const options = { + options: {}, + context: createMockContext(), + }; + + wrapMethodWithSentry(options, handler, undefined, true); + + expect(isInstrumented(handler)).toBeFalsy(); + }); + }); + + describe('span creation', () => { + it('creates span with spanName when provided', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context: createMockContext(), + spanName: 'test-span', + spanOp: 'test-op', + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(sentryCore.startSpan).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'test-span', + }), + expect.any(Function), + ); + }); + + it('does not create span when spanName is not provided', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + // startSpan should not be called when no spanName is provided + expect(sentryCore.startSpan).not.toHaveBeenCalled(); + }); + }); + + describe('error handling', () => { + it('captures exceptions from sync methods', async () => { + const error = new Error('Test sync error'); + const handler = vi.fn().mockImplementation(() => { + throw error; + }); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + await expect(async () => wrapped()).rejects.toThrow('Test sync error'); + expect(sentryCore.captureException).toHaveBeenCalledWith(error, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + }); + + it('captures exceptions from async methods', async () => { + const error = new Error('Test async error'); + const handler = vi.fn().mockRejectedValue(error); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + await expect(wrapped()).rejects.toThrow('Test async error'); + expect(sentryCore.captureException).toHaveBeenCalledWith(error, { + mechanism: { + type: 'auto.faas.cloudflare.durable_object', + handled: false, + }, + }); + }); + }); + + describe('callback execution', () => { + it('executes callback before handler', async () => { + const callOrder: string[] = []; + const handler = vi.fn().mockImplementation(() => { + callOrder.push('handler'); + return 'result'; + }); + const callback = vi.fn().mockImplementation(() => { + callOrder.push('callback'); + }); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler, callback); + await wrapped('arg1', 'arg2'); + + expect(callback).toHaveBeenCalledWith('arg1', 'arg2'); + expect(callOrder).toEqual(['callback', 'handler']); + }); + }); + + describe('waitUntil flush', () => { + it('calls waitUntil with flush when context has waitUntil', async () => { + const waitUntil = vi.fn(); + const context = { + waitUntil, + originalStorage: undefined, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context, + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped(); + + expect(waitUntil).toHaveBeenCalled(); + expect(sentryCore.flush).toHaveBeenCalledWith(2000); + }); + + it('handles missing waitUntil gracefully', async () => { + const context = { + originalStorage: undefined, + } as any; + + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context, + }; + + const wrapped = wrapMethodWithSentry(options, handler); + + // Should not throw + await expect(wrapped()).resolves.toBeDefined(); + }); + }); + + describe('argument passing', () => { + it('passes arguments to handler', async () => { + const handler = vi.fn().mockResolvedValue('result'); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped('arg1', 'arg2', { key: 'value' }); + + expect(handler).toHaveBeenCalledWith('arg1', 'arg2', { key: 'value' }); + }); + + it('preserves this context', async () => { + const thisArg = { name: 'test-context' }; + const handler = vi.fn(function (this: any) { + return this.name; + }); + const options = { + options: {}, + context: createMockContext(), + }; + + const wrapped = wrapMethodWithSentry(options, handler); + await wrapped.call(thisArg); + + expect(handler.mock.instances[0]).toBe(thisArg); + }); + }); +});