observability-test: isolate ObservabilityOptions injection tests

googleapis · Oct 4, 2024 · 692a62c · 692a62c
1 parent b69e3b9
commit 692a62c
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 48 deletions.
diff --git a/observability-test/spanner.ts b/observability-test/spanner.ts
@@ -206,6 +206,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.getSnapshot',
  ];
@@ -249,6 +250,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.getTransaction',
  ];
@@ -293,6 +295,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.runStream',
  ];
@@ -339,6 +342,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.runStream',
  'CloudSpanner.Database.run',
@@ -401,6 +405,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.runTransaction',
  ];
@@ -445,6 +450,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.writeAtLeastOnce',
  ];
@@ -488,6 +494,7 @@ describe('EndToEnd', () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.batchCreateSessions',
  ];
@@ -515,46 +522,32 @@ describe('EndToEnd', () => {
 });
 
 describe('ObservabilityOptions injection and propagation', async () => {
- const globalTraceExporter = new InMemorySpanExporter();
- const globalTracerProvider = new NodeTracerProvider({
- sampler: new AlwaysOnSampler(),
- exporter: globalTraceExporter,
- });
- globalTracerProvider.addSpanProcessor(
- new SimpleSpanProcessor(globalTraceExporter)
- );
- globalTracerProvider.register();
-
- const injectedTraceExporter = new InMemorySpanExporter();
- const injectedTracerProvider = new NodeTracerProvider({
- sampler: new AlwaysOnSampler(),
- exporter: injectedTraceExporter,
- });
- injectedTracerProvider.addSpanProcessor(
- new SimpleSpanProcessor(injectedTraceExporter)
- );
-
- const observabilityOptions: typeof ObservabilityOptions = {
- tracerProvider: injectedTracerProvider,
- enableExtendedTracing: true,
- };
-
- const setupResult = await setup(observabilityOptions);
- const spanner = setupResult.spanner;
- const server = setupResult.server;
- const spannerMock = setupResult.spannerMock;
-
- after(async () => {
- globalTraceExporter.reset();
- injectedTraceExporter.reset();
- await globalTracerProvider.shutdown();
- await injectedTracerProvider.shutdown();
- spannerMock.resetRequests();
- spanner.close();
- server.tryShutdown(() => {});
- });
+ it('Passed into Spanner, Instance and Database', async () => {
+ const traceExporter = new InMemorySpanExporter();
+ const tracerProvider = new NodeTracerProvider({
+ sampler: new AlwaysOnSampler(),
+ exporter: traceExporter,
+ });
+ tracerProvider.addSpanProcessor(new SimpleSpanProcessor(traceExporter));
+
+ const observabilityOptions: typeof ObservabilityOptions = {
+ tracerProvider: tracerProvider,
+ enableExtendedTracing: true,
+ };
+
+ const setupResult = await setup(observabilityOptions);
+ const spanner = setupResult.spanner;
+ const server = setupResult.server;
+ const spannerMock = setupResult.spannerMock;
+
+ after(async () => {
+ traceExporter.reset();
+ await tracerProvider.shutdown();
+ spannerMock.resetRequests();
+ spanner.close();
+ server.tryShutdown(() => {});
+ });
 
- it('Passed into Spanner, Instance and Database', done => {
  // Ensure that the same observability configuration is set on the Spanner client.
  assert.deepStrictEqual(spanner._observabilityOptions, observabilityOptions);
 
@@ -588,11 +581,45 @@ describe('ObservabilityOptions injection and propagation', async () => {
  databaseByConstructor._observabilityOptions,
  observabilityOptions
  );
-
- done();
  });
 
- it('Propagates spans to the injected not global TracerProvider', done => {
+ it('Propagates spans to the injected not global TracerProvider', async () => {
+ const globalTraceExporter = new InMemorySpanExporter();
+ const globalTracerProvider = new NodeTracerProvider({
+ sampler: new AlwaysOnSampler(),
+ exporter: globalTraceExporter,
+ });
+ globalTracerProvider.addSpanProcessor(
+ new SimpleSpanProcessor(globalTraceExporter)
+ );
+ globalTracerProvider.register();
+
+ const injectedTraceExporter = new InMemorySpanExporter();
+ const injectedTracerProvider = new NodeTracerProvider({
+ sampler: new AlwaysOnSampler(),
+ exporter: injectedTraceExporter,
+ });
+ injectedTracerProvider.addSpanProcessor(
+ new SimpleSpanProcessor(injectedTraceExporter)
+ );
+
+ const observabilityOptions: typeof ObservabilityOptions = {
+ tracerProvider: injectedTracerProvider,
+ enableExtendedTracing: true,
+ };
+ const setupResult = await setup(observabilityOptions);
+ const spanner = setupResult.spanner;
+ const server = setupResult.server;
+ const spannerMock = setupResult.spannerMock;
+
+ after(async () => {
+ injectedTraceExporter.reset();
+ await injectedTracerProvider.shutdown();
+ spannerMock.resetRequests();
+ spanner.close();
+ server.tryShutdown(() => {});
+ });
+
  const instance = spanner.instance('instance');
  const database = instance.database('database');
 
@@ -628,6 +655,7 @@ describe('ObservabilityOptions injection and propagation', async () => {
  });
 
  const expectedSpanNames = [
+ 'CloudSpanner.Database.batchCreateSessions',
  'CloudSpanner.SessionPool.createSessions',
  'CloudSpanner.Database.runStream',
  'CloudSpanner.Database.run',
@@ -652,8 +680,6 @@ describe('ObservabilityOptions injection and propagation', async () => {
  expectedEventNames,
  `Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
  );
-
- done();
  });
  });
 });
diff --git a/src/database.ts b/src/database.ts
@@ -450,14 +450,14 @@ class Database extends common.GrpcServiceObject {
  typeof poolOptions === 'function'
  ? new (poolOptions as SessionPoolConstructor)(this, null)
  : new SessionPool(this, poolOptions);
- const concretePool = this.pool_ as SessionPool;
- if (concretePool) {
- concretePool._observabilityOptions = instance._observabilityOptions;
+ const sessionPoolInstance = this.pool_ as SessionPool;
+ if (sessionPoolInstance) {
+ sessionPoolInstance._observabilityOptions =
+ instance._observabilityOptions;
  }
  if (typeof poolOptions === 'object') {
  this.databaseRole = poolOptions.databaseRole || null;
  }
- this._observabilityOptions = instance._observabilityOptions;
  this.formattedName_ = formattedName_;
  this.instance = instance;
  this.resourceHeader_ = {
@@ -472,6 +472,7 @@ class Database extends common.GrpcServiceObject {
  Object.assign({}, queryOptions),
  Database.getEnvironmentQueryOptions()
  );
+ this._observabilityOptions = instance._observabilityOptions;
  }
  /**
  * @typedef {array} SetDatabaseMetadataResponse