Add callback variant of regression test

googleapis · Oct 9, 2024 · bb44c62 · bb44c62
1 parent 6c879ec
commit bb44c62
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 40 deletions.
diff --git a/observability-test/database.ts b/observability-test/database.ts
@@ -518,7 +518,7 @@ describe('Database', () => {
 
       getSessionStub.callsFake(callback => callback(fakeError, null));
 
-      database.getSnapshot((err, snapshot) => {
+      database.getSnapshot(err => {
         assert.strictEqual(err, fakeError);
         traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
@@ -1136,7 +1136,6 @@ describe('Database', () => {
     });
 
     it('with error on null mutation should catch thrown error', done => {
-      const fakeError = new Error('err');
       try {
         database.writeAtLeastOnce(null, (err, res) => {});
       } catch (err) {

diff --git a/observability-test/session-pool.ts b/observability-test/session-pool.ts
@@ -32,7 +32,6 @@ const {SimpleSpanProcessor} = require('@opentelemetry/sdk-trace-base');
 import {Database} from '../src/database';
 import {Session} from '../src/session';
 import * as sp from '../src/session-pool';
-import {Transaction} from '../src/transaction';
 
 let pQueueOverride: typeof PQueue | null = null;
 

diff --git a/observability-test/spanner.ts b/observability-test/spanner.ts
@@ -1010,52 +1010,43 @@ describe('ObservabilityOptions injection and propagation', async () => {
   });
 });
 
-describe('Regression tests for fixed bugs', () => {
-  it('async/await correctly parents trace spans', async () => {
-    // See https://github.com/googleapis/nodejs-spanner/issues/2146.
-    const traceExporter = new InMemorySpanExporter();
-    const provider = new NodeTracerProvider({
+describe('Regression tests for fixed bugs', async () => {
+  let server: grpc.Server;
+  let spanner: Spanner;
+  let database: Database;
+  let spannerMock: mock.MockSpanner;
+  let traceExporter: typeof InMemorySpanExporter;
+  let provider: typeof TracerProvider;
+  let observabilityOptions: typeof ObservabilityOptions;
+
+  beforeEach(async () => {
+    traceExporter = new InMemorySpanExporter();
+    provider = new NodeTracerProvider({
       sampler: new AlwaysOnSampler(),
       exporter: traceExporter,
     });
     provider.addSpanProcessor(new SimpleSpanProcessor(traceExporter));
 
-    const observabilityOptions: typeof ObservabilityOptions = {
+    observabilityOptions = {
       tracerProvider: provider,
       enableExtendedTracing: true,
     };
     const setupResult = await setup(observabilityOptions);
-    const spanner = setupResult.spanner;
-    const server = setupResult.server;
-    const spannerMock = setupResult.spannerMock;
-
-    after(async () => {
-      provider.shutdown();
-      spannerMock.resetRequests();
-      spanner.close();
-      server.tryShutdown(() => {});
-    });
-
-    async function main() {
-      const instance = spanner.instance('testing');
-      instance._observabilityOptions = observabilityOptions;
-      const database = instance.database('db-1');
-
-      const query = {
-        sql: selectSql,
-      };
-
-      const [rows] = await database.run(query);
-
-      rows.forEach(row => {
-        const json = row.toJSON();
-      });
-
-      provider.forceFlush();
-    }
+    spanner = setupResult.spanner;
+    server = setupResult.server;
+    spannerMock = setupResult.spannerMock;
+  });
 
-    await main();
+  afterEach(async () => {
+    traceExporter.reset();
+    provider.shutdown();
+    spannerMock.resetRequests();
+    spanner.close();
+    server.tryShutdown(() => {});
+  });
 
+  function assertAsyncAwaitExpectations() {
+    // See https://github.com/googleapis/nodejs-spanner/issues/2146.
     traceExporter.forceFlush();
     const spans = traceExporter.getFinishedSpans();
 
@@ -1162,5 +1153,55 @@ describe('Regression tests for fixed bugs', () => {
       expectedEventNames,
       `Unexpected events:\n\tGot:  ${actualEventNames}\n\tWant: ${expectedEventNames}`
     );
+  }
+
+  it('async/await correctly parents trace spans', async () => {
+    // See https://github.com/googleapis/nodejs-spanner/issues/2146.
+    async function main() {
+      const instance = spanner.instance('testing');
+      instance._observabilityOptions = observabilityOptions;
+      const database = instance.database('db-1');
+
+      const query = {
+        sql: selectSql,
+      };
+
+      const [rows] = await database.run(query);
+
+      rows.forEach(row => {
+        const json = row.toJSON();
+      });
+
+      provider.forceFlush();
+    }
+
+    await main();
+    assertAsyncAwaitExpectations();
+  });
+
+  it('callback correctly parents trace spans', done => {
+    function main(onComplete) {
+      const instance = spanner.instance('testing');
+      instance._observabilityOptions = observabilityOptions;
+      const database = instance.database('db-1');
+
+      const query = {
+        sql: selectSql,
+      };
+
+      database.run(query, (err, rows) => {
+        rows.forEach(row => {
+          const json = row.toJSON();
+        });
+
+        provider.forceFlush();
+        onComplete();
+      });
+    }
+
+    main(() => {
+      assertAsyncAwaitExpectations();
+      done();
+    });
   });
 });
diff --git a/src/session-pool.ts b/src/session-pool.ts
@@ -762,8 +762,8 @@ export class SessionPool extends EventEmitter implements SessionPoolInterface {
       this.database._observabilityOptions = this._observabilityOptions;
     }
 
-    const q = {opts: this._observabilityOptions};
-    return startTrace('SessionPool.createSessions', q, async span => {
+    const traceConfig = {opts: this._observabilityOptions};
+    return startTrace('SessionPool.createSessions', traceConfig, async span => {
       span.addEvent(`Requesting ${amount} sessions`);
 
       // while we can request as many sessions be created as we want, the backend