Add callback variant of regression test

Author: odeke-em

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) {

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;
 

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();
+ });
  });
 });

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