test: refactor

observability-test/database.ts

@@ -47,6 +47,7 @@ import {Instance, MutationGroup, Spanner} from '../src';
 import * as pfy from '@google-cloud/promisify';
 import {grpc} from 'google-gax';
 import {MockError} from '../test/mockserver/mockspanner';
+import {FakeSessionFactory} from '../test/database';
 const {generateWithAllSpansHaveDBName} = require('./helper');
 
 const fakePfy = extend({}, pfy, {
@@ -234,6 +235,7 @@ describe('Database', () => {
       './codec': {codec: fakeCodec},
       './partial-result-stream': {partialResultStream: fakePartialResultStream},
       './session-pool': {SessionPool: FakeSessionPool},
+      './session-factory': {SessionFactory: FakeSessionFactory},
       './session': {Session: FakeSession},
       './table': {Table: FakeTable},
       './transaction-runner': {

observability-test/spanner.ts

@@ -505,7 +505,7 @@ describe('ObservabilityOptions injection and propagation', async () => {
       db.formattedName_
     );
 
-    it('run', () => {
+    it('run', done => {
       database.getTransaction((err, tx) => {
         assert.ifError(err);
 
@@ -549,6 +549,8 @@ describe('ObservabilityOptions injection and propagation', async () => {
             true,
             `Unexpected events:\n\tGot:  ${actualEventNames}\n\tWant: ${expectedEventNames}`
           );
+
+          done();
         });
       });
     });
@@ -608,16 +610,14 @@ describe('ObservabilityOptions injection and propagation', async () => {
       });
     });
 
-    it('runStream', () => {
+    it('runStream', done => {
       let rowCount = 0;
       database.getTransaction((err, tx) => {
         assert.ifError(err);
         tx!
           .runStream(selectSql)
           .on('data', () => rowCount++)
-          .on('error', () => {
-            assert.ifError;
-          })
+          .on('error', assert.ifError)
           .on('stats', () => {})
           .on('end', async () => {
             tx!.end();
@@ -657,6 +657,8 @@ describe('ObservabilityOptions injection and propagation', async () => {
               expectedEventNames,
               `Unexpected events:\n\tGot:  ${actualEventNames}\n\tWant: ${expectedEventNames}`
             );
+
+            done();
           });
       });
     });

src/database.ts

@@ -36,6 +36,7 @@ import {
 } from 'google-gax';
 import {Backup} from './backup';
 import {BatchTransaction, TransactionIdentifier} from './batch-transaction';
+import {SessionFactory, SessionFactoryInterface} from './session-factory';
 import {
   google as databaseAdmin,
   google,
@@ -102,7 +103,6 @@ import Policy = google.iam.v1.Policy;
 import FieldMask = google.protobuf.FieldMask;
 import IDatabase = google.spanner.admin.database.v1.IDatabase;
 import snakeCase = require('lodash.snakecase');
-import {SessionFactory, SessionFactoryInterface} from './session-factory';
 import {
   ObservabilityOptions,
   Span,
@@ -458,6 +458,7 @@ class Database extends common.GrpcServiceObject {
     }
     this.formattedName_ = formattedName_;
     this.instance = instance;
+    this._observabilityOptions = instance._observabilityOptions;
     this._traceConfig = {
       opts: this._observabilityOptions,
       dbName: this.formattedName_,
@@ -472,7 +473,6 @@ class Database extends common.GrpcServiceObject {
     this.requestStream = instance.requestStream as any;
     this.sessionFactory_ = new SessionFactory(this, name, poolOptions);
     this.pool_ = this.sessionFactory_.getPool();
-    this.multiplexedSession_ = this.sessionFactory_.getMultiplexedSession();
     const sessionPoolInstance = this.pool_ as SessionPool;
     if (sessionPoolInstance) {
       sessionPoolInstance._observabilityOptions =

src/session-factory.ts

@@ -27,7 +27,6 @@ import {
 import {SessionPoolConstructor} from './database';
 import {ServiceObjectConfig} from '@google-cloud/common';
 const common = require('./common-grpc/service-object');
-
 /**
  * @callback GetSessionCallback
  * @param {?Error} error Request error, if any.

src/session-pool.ts

@@ -24,14 +24,13 @@ import {Transaction} from './transaction';
 import {NormalCallback} from './common';
 import {GoogleError, grpc, ServiceError} from 'google-gax';
 import trace = require('stack-trace');
-import {GetSessionCallback} from './session-factory';
 import {
   ObservabilityOptions,
   getActiveOrNoopSpan,
   setSpanErrorAndException,
   startTrace,
 } from './instrument';
-
+import {GetSessionCallback} from './session-factory';
 import {
   isDatabaseNotFoundError,
   isInstanceNotFoundError,

test/database.ts

@@ -91,7 +91,7 @@ function fakePartialResultStream(this: Function & {calledWith_: IArguments}) {
   return this;
 }
 
-class FakeSession {
+export class FakeSession {
   calledWith_: IArguments;
   formattedName_: any;
   constructor() {
@@ -109,42 +109,46 @@ class FakeSession {
   }
 }
 
-export class FakeSessionFactory extends EventEmitter {
+export class FakeSessionPool extends EventEmitter {
   calledWith_: IArguments;
   constructor() {
     super();
     this.calledWith_ = arguments;
   }
-  getSession(): FakeSession {
-    return new FakeSession();
-  }
-  getPool(): FakeSessionPool {
-    return new FakeSessionPool();
-  }
-  getMultiplexedSession(): FakeMultiplexedSession {
-    return new FakeMultiplexedSession();
-  }
+  open() {}
+  getSession() {}
+  release() {}
 }
 
-export class FakeSessionPool extends EventEmitter {
+export class FakeMultiplexedSession extends EventEmitter {
   calledWith_: IArguments;
   constructor() {
     super();
     this.calledWith_ = arguments;
   }
-  open() {}
+  createSession() {}
   getSession() {}
-  release() {}
 }
 
-export class FakeMultiplexedSession extends EventEmitter {
+export class FakeSessionFactory extends EventEmitter {
   calledWith_: IArguments;
   constructor() {
     super();
     this.calledWith_ = arguments;
   }
-  createSession() {}
-  getSession() {}
+  getSession(): FakeSession | FakeMultiplexedSession {
+    if (process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS) {
+      return new FakeSession();
+    } else {
+      return new FakeMultiplexedSession();
+    }
+  }
+  getPool(): FakeSessionPool {
+    return new FakeSessionPool();
+  }
+  getMultiplexedSession(): FakeMultiplexedSession {
+    return new FakeMultiplexedSession();
+  }
 }
 
 class FakeTable {
@@ -269,8 +273,9 @@ describe('Database', () => {
       './batch-transaction': {BatchTransaction: FakeBatchTransaction},
       './codec': {codec: fakeCodec},
       './partial-result-stream': {partialResultStream: fakePartialResultStream},
-      './session': {Session: FakeSession},
+      './session-pool': {SessionPool: FakeSessionPool},
       './session-factory': {SessionFactory: FakeSessionFactory},
+      './session': {Session: FakeSession},
       './table': {Table: FakeTable},
       './transaction-runner': {
         TransactionRunner: FakeTransactionRunner,