From 692a62c37874264a737b9f2013756447e0db6f74 Mon Sep 17 00:00:00 2001
From: Emmanuel T Odeke <emmanuel@orijtech.com>
Date: Fri, 4 Oct 2024 00:31:43 -0700
Subject: [PATCH] observability-test: isolate ObservabilityOptions injection
 tests

---
 observability-test/spanner.ts | 114 +++++++++++++++++++++-------------
 src/database.ts               |   9 +--
 2 files changed, 75 insertions(+), 48 deletions(-)

diff --git a/observability-test/spanner.ts b/observability-test/spanner.ts
index d308898fd..977c7943a 100644
--- 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
index a756f77c5..168ff34f4 100644
--- 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