ref: Rewrite grpc test to be less accurate, but more stable

AbhiPrasad · AbhiPrasad · commit a12a49326509 · 2025-03-25T18:48:31.000-04:00
This rewrite makes the tests less accurate because we are no longer
testing against google-gax directly, but this feels like something we
should be doing in our integration tests instead.

As an alternative we check that we are filling the prototype correctly,
and then if we are patching the google-gax stub properly. This also
allows us to drop googleapis and google-gax as dev deps.
diff --git a/packages/google-cloud-serverless/package.json b/packages/google-cloud-serverless/package.json
@@ -58,7 +58,6 @@
     "@google-cloud/functions-framework": "^1.7.1",
     "@google-cloud/pubsub": "^2.5.0",
     "@types/node": "^18.19.1",
-    "google-gax": "^2.9.0",
     "nock": "^13.5.5"
   },
   "scripts": {
diff --git a/packages/google-cloud-serverless/src/integrations/google-cloud-grpc.ts b/packages/google-cloud-serverless/src/integrations/google-cloud-grpc.ts
@@ -3,11 +3,11 @@ import type { Client, IntegrationFn } from '@sentry/core';
 import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration, fill, getClient } from '@sentry/core';
 import { startInactiveSpan } from '@sentry/node';
 
-interface GrpcFunction extends CallableFunction {
+export interface GrpcFunction extends CallableFunction {
   (...args: unknown[]): EventEmitter;
 }
 
-interface GrpcFunctionObject extends GrpcFunction {
+export interface GrpcFunctionObject extends GrpcFunction {
   requestStream: boolean;
   responseStream: boolean;
   originalName: string;
@@ -21,7 +21,7 @@ interface CreateStubFunc extends CallableFunction {
   (createStub: unknown, options: StubOptions): PromiseLike<Stub>;
 }
 
-interface Stub {
+export interface Stub {
   [key: string]: GrpcFunctionObject;
 }
 
@@ -78,7 +78,7 @@ function wrapCreateStub(origCreate: CreateStubFunc): CreateStubFunc {
 }
 
 /** Patches the function in grpc stub to enable tracing */
-function fillGrpcFunction(stub: Stub, serviceIdentifier: string, methodName: string): void {
+export function fillGrpcFunction(stub: Stub, serviceIdentifier: string, methodName: string): void {
   const funcObj = stub[methodName];
   if (typeof funcObj !== 'function') {
     return;
diff --git a/packages/google-cloud-serverless/test/integrations/google-cloud-grpc.test.ts b/packages/google-cloud-serverless/test/integrations/google-cloud-grpc.test.ts
@@ -1,28 +1,30 @@
-import * as dns from 'dns';
-import { EventEmitter } from 'events';
-import * as fs from 'fs';
-import * as path from 'path';
-import { PubSub } from '@google-cloud/pubsub';
-import * as http2 from 'http2';
-// eslint-disable-next-line @typescript-eslint/ban-ts-comment
-// @ts-ignore ESM/CJS interop issue
-import nock from 'nock';
-import { describe, vi, beforeEach, test, expect, type Mock, afterAll, afterEach } from 'vitest';
-
-import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
-import { NodeClient, createTransport, setCurrentClient } from '@sentry/node';
-import { googleCloudGrpcIntegration } from '../../src/integrations/google-cloud-grpc';
-
-vi.mock('dns');
-vi.mock('http2');
-
-const spyConnect = vi.spyOn(http2, 'connect');
+import { vi, describe, beforeEach, test, expect } from 'vitest';
+import { NodeClient } from '@sentry/node';
+import { createTransport } from '@sentry/core';
+import { setCurrentClient, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
+import { googleCloudGrpcIntegration, fillGrpcFunction } from '../../src/integrations/google-cloud-grpc';
+import type { GrpcFunctionObject, Stub, GrpcFunction } from '../../src/integrations/google-cloud-grpc';
+
 const mockSpanEnd = vi.fn();
-const mockStartInactiveSpan = vi.fn(spanArgs => ({ ...spanArgs }));
+const mockStartInactiveSpan = vi.fn();
+const mockFill = vi.fn();
+
+let mockClient: NodeClient;
+
+vi.mock('@sentry/core', async () => {
+  const original = await vi.importActual('@sentry/core');
+  return {
+    ...original,
+    fill: (obj: any, name: string, replacement: any) => {
+      mockFill(obj, name, replacement);
+      obj[name] = replacement(obj[name]);
+    },
+    getClient: () => mockClient,
+  };
+});
 
 vi.mock('@sentry/node', async () => {
-  // eslint-disable-next-line @typescript-eslint/consistent-type-imports
-  const original = (await vi.importActual('@sentry/node')) as typeof import('@sentry/node');
+  const original = await vi.importActual('@sentry/node');
   return {
     ...original,
     startInactiveSpan: (ctx: unknown) => {
@@ -32,130 +34,177 @@ vi.mock('@sentry/node', async () => {
   };
 });
 
-/** Fake HTTP2 stream */
-class FakeStream extends EventEmitter {
-  public rstCode: number = 0;
-  close() {
-    this.emit('end');
-    this.emit('close');
-  }
-  end() {}
-  pause() {}
-  resume() {}
-  write(_data: Buffer, cb: CallableFunction) {
-    process.nextTick(cb, null);
-  }
-}
-
-/** Fake HTTP2 session for GRPC */
-class FakeSession extends EventEmitter {
-  public socket: EventEmitter = new EventEmitter();
-  public request: Mock = vi.fn();
-  ping() {}
-  mockRequest(fn: (stream: FakeStream) => void): FakeStream {
-    const stream = new FakeStream();
-    this.request.mockImplementationOnce(() => {
-      process.nextTick(fn, stream);
-      return stream;
-    });
-    return stream;
-  }
-  mockUnaryRequest(responseData: Buffer) {
-    this.mockRequest(stream => {
-      stream.emit(
-        'response',
-        { ':status': 200, 'content-type': 'application/grpc', 'content-disposition': 'attachment' },
-        4,
-      );
-      stream.emit('data', responseData);
-      stream.emit('trailers', { 'grpc-status': '0', 'content-disposition': 'attachment' });
-    });
-  }
-  close() {
-    this.emit('close');
-    this.socket.emit('close');
-  }
-  ref() {}
-  unref() {}
+// Need to override mock because the integration loads google-gax as a CJS file
+async function mock(mockedUri: string, stub: any) {
+  // @ts-expect-error we are using import on purpose
+  const { Module } = await import('module');
+
+  // @ts-expect-error test
+  Module._load_original = Module._load;
+  // @ts-expect-error test
+  Module._load = (uri, parent) => {
+    if (uri === mockedUri) return stub;
+    // @ts-expect-error test
+    return Module._load_original(uri, parent);
+  };
 }
 
-function mockHttp2Session(): FakeSession {
-  const session = new FakeSession();
-  spyConnect.mockImplementationOnce(() => {
-    process.nextTick(() => session.emit('connect'));
-    return session as unknown as http2.ClientHttp2Session;
-  });
-  return session;
-}
+vi.hoisted(
+  () =>
+    void mock('google-gax', {
+      GrpcClient: {
+        prototype: {
+          createStub: vi.fn(),
+        },
+      },
+    }),
+);
 
 describe('GoogleCloudGrpc tracing', () => {
-  const mockClient = new NodeClient({
-    tracesSampleRate: 1.0,
-    integrations: [],
-    dsn: 'https://withAWSServices@domain/123',
-    transport: () => createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
-    stackParser: () => [],
-  });
+  beforeEach(() => {
+    mockClient = new NodeClient({
+      tracesSampleRate: 1.0,
+      integrations: [],
+      dsn: 'https://withAWSServices@domain/123',
+      transport: () => createTransport({ recordDroppedEvent: () => undefined }, _ => Promise.resolve({})),
+      stackParser: () => [],
+    });
 
-  const integration = googleCloudGrpcIntegration();
-  mockClient.addIntegration(integration);
+    const integration = googleCloudGrpcIntegration();
+    mockClient.addIntegration(integration);
+    integration.setup?.(mockClient);
 
-  beforeEach(() => {
-    nock('https://www.googleapis.com').post('/oauth2/v4/token').reply(200, []);
     setCurrentClient(mockClient);
     mockSpanEnd.mockClear();
     mockStartInactiveSpan.mockClear();
+    mockFill.mockClear();
   });
 
-  afterAll(() => {
-    nock.restore();
-    spyConnect.mockRestore();
-  });
-
-  // We use google cloud pubsub as an example of grpc service for which we can trace requests.
-  describe('pubsub', () => {
-    // @ts-expect-error see "Why @ts-expect-error" note
-    const dnsLookup = dns.lookup as Mock;
-    // @ts-expect-error see "Why @ts-expect-error" note
-    const resolveTxt = dns.resolveTxt as Mock;
-    dnsLookup.mockImplementation((hostname, ...args) => {
-      expect(hostname).toEqual('pubsub.googleapis.com');
-      process.nextTick(args[args.length - 1], null, [{ address: '0.0.0.0', family: 4 }]);
-    });
-    resolveTxt.mockImplementation((hostname, cb) => {
-      expect(hostname).toEqual('pubsub.googleapis.com');
-      process.nextTick(cb, null, []);
+  describe('setup', () => {
+    test('integration name is correct', () => {
+      const integration = googleCloudGrpcIntegration();
+      expect(integration.name).toBe('GoogleCloudGrpc');
     });
 
-    const pubsub = new PubSub({
-      credentials: {
-        client_email: 'client@email',
-        private_key: fs.readFileSync(path.resolve(__dirname, 'private.pem')).toString(),
-      },
-      projectId: 'project-id',
+    test('setupOnce patches GrpcClient.createStub', () => {
+      const mockCreateStub = vi.fn();
+      const mockGrpcClient = {
+        prototype: {
+          createStub: mockCreateStub,
+        },
+      };
+
+      // eslint-disable-next-line @typescript-eslint/no-var-requires
+      require('google-gax').GrpcClient = mockGrpcClient;
+
+      const integration = googleCloudGrpcIntegration();
+      integration.setupOnce?.();
+      expect(mockCreateStub).toBeDefined();
     });
 
-    afterEach(() => {
-      dnsLookup.mockReset();
-      resolveTxt.mockReset();
+    test('setupOnce throws when google-gax is not available and not optional', () => {
+      // eslint-disable-next-line @typescript-eslint/no-var-requires
+      require('google-gax').GrpcClient = undefined;
+
+      const integration = googleCloudGrpcIntegration();
+      expect(() => integration.setupOnce?.()).toThrow();
     });
 
-    afterAll(async () => {
-      await pubsub.close();
+    test('setupOnce does not throw when google-gax is not available and optional', () => {
+      // eslint-disable-next-line @typescript-eslint/no-var-requires
+      require('google-gax').GrpcClient = undefined;
+
+      const optionalIntegration = googleCloudGrpcIntegration({ optional: true });
+      expect(() => optionalIntegration.setupOnce?.()).not.toThrow();
     });
+  });
 
-    test('publish', async () => {
-      mockHttp2Session().mockUnaryRequest(Buffer.from('00000000120a1031363337303834313536363233383630', 'hex'));
-      const resp = await pubsub.topic('nicetopic').publish(Buffer.from('data'));
-      expect(resp).toEqual('1637084156623860');
-      expect(mockStartInactiveSpan).toBeCalledWith({
-        op: 'grpc.pubsub',
+  describe('fillGrpcFunction', () => {
+    test('patches unary call methods with tracing', () => {
+      const mockStub: Stub = {
+        unaryMethod: Object.assign(vi.fn(), {
+          requestStream: false,
+          responseStream: false,
+          originalName: 'unaryMethod',
+        } as GrpcFunctionObject),
+      };
+
+      const mockEventEmitter = {
+        on: vi.fn(),
+      };
+
+      (mockStub.unaryMethod as any).apply = vi.fn().mockReturnValue(mockEventEmitter);
+
+      fillGrpcFunction(mockStub, 'test-service', 'unaryMethod');
+
+      const result = (mockStub.unaryMethod as GrpcFunction)();
+      expect(result).toBe(mockEventEmitter);
+      expect(mockEventEmitter.on).toHaveBeenCalledWith('status', expect.any(Function));
+      expect(mockStartInactiveSpan).toHaveBeenCalledWith({
+        name: 'unary call unaryMethod',
+        onlyIfParent: true,
+        op: 'grpc.test-service',
         attributes: {
           [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.grpc.serverless',
         },
-        name: 'unary call publish',
-        onlyIfParent: true,
       });
     });
+
+    test('does not patch non-unary call methods', () => {
+      const mockStub: Stub = {
+        clientStreamMethod: Object.assign(vi.fn(), {
+          requestStream: true,
+          responseStream: false,
+          originalName: 'clientStreamMethod',
+        } as GrpcFunctionObject),
+        serverStreamMethod: Object.assign(vi.fn(), {
+          requestStream: false,
+          responseStream: true,
+          originalName: 'serverStreamMethod',
+        } as GrpcFunctionObject),
+        bidiStreamMethod: Object.assign(vi.fn(), {
+          requestStream: true,
+          responseStream: true,
+          originalName: 'bidiStreamMethod',
+        } as GrpcFunctionObject),
+      };
+
+      fillGrpcFunction(mockStub, 'test-service', 'clientStreamMethod');
+      fillGrpcFunction(mockStub, 'test-service', 'serverStreamMethod');
+      fillGrpcFunction(mockStub, 'test-service', 'bidiStreamMethod');
+
+      expect(mockStartInactiveSpan).not.toHaveBeenCalled();
+    });
+
+    test('does not patch non-function properties', () => {
+      const mockStub: Stub = {
+        nonFunction: Object.assign(vi.fn(), {
+          requestStream: false,
+          responseStream: false,
+          originalName: 'nonFunction',
+        } as GrpcFunctionObject),
+      };
+
+      fillGrpcFunction(mockStub, 'test-service', 'nonFunction');
+      expect(mockStartInactiveSpan).not.toHaveBeenCalled();
+    });
+
+    test('does not patch methods when return value is not an EventEmitter', () => {
+      const mockStub: Stub = {
+        unaryMethod: Object.assign(vi.fn(), {
+          requestStream: false,
+          responseStream: false,
+          originalName: 'unaryMethod',
+        } as GrpcFunctionObject),
+      };
+
+      (mockStub.unaryMethod as any).apply = vi.fn().mockReturnValue({ notAnEventEmitter: true });
+
+      fillGrpcFunction(mockStub, 'test-service', 'unaryMethod');
+
+      const result = (mockStub.unaryMethod as GrpcFunction)();
+      expect(result).toEqual({ notAnEventEmitter: true });
+      expect(mockStartInactiveSpan).not.toHaveBeenCalled();
+    });
   });
 });
diff --git a/yarn.lock b/yarn.lock
@@ -17499,7 +17499,7 @@ google-auth-library@^7.0.0, google-auth-library@^7.0.2:
     jws "^4.0.0"
     lru-cache "^6.0.0"
 
-google-gax@^2.9.0, google-gax@^2.9.2:
+google-gax@^2.9.2:
   version "2.11.2"
   resolved "https://registry.yarnpkg.com/google-gax/-/google-gax-2.11.2.tgz#9ef7773b94aaa61c4588fb2408d62e8444995026"
   integrity sha512-PNqXv7Oi5XBMgoMWVxLZHUidfMv7cPHrDSDXqLyEd6kY6pqFnVKC8jt2T1df4JPSc2+VLPdeo6L7X9mbdQG8Xw==
