chore: refactor observability tests (#2177)

* chore: integration test fix * traces tests refactoring * feat: (observability): trace Database.runPartitionedUpdate (#2176) This change traces Database.runPartitionedUpdate along with the appropriate tests for it with and without errors. Updates #2079 * moving additional attributes to separate PR --------- Co-authored-by: Emmanuel T Odeke <[email protected]> Co-authored-by: alkatrivedi <[email protected]>
googleapis · Oct 30, 2024 · f08cf4b · f08cf4b
1 parent be063b0
commit f08cf4b
Show file tree

Hide file tree

Showing 12 changed files with 263 additions and 671 deletions.
diff --git a/observability-test/batch-transaction.ts b/observability-test/batch-transaction.ts
@@ -153,103 +153,95 @@ describe('BatchTransaction', () => {
   };
 
   it('createQueryPartitions', done => {
-    const REQUEST = sandbox.stub();
-
-    const res = batchTransaction.createQueryPartitions(
-      QUERY,
-      (err, part, resp) => {
-        assert.ifError(err);
-        traceExporter.forceFlush();
-        const spans = traceExporter.getFinishedSpans();
-        assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
-
-        // Sort the spans by duration.
-        spans.sort((spanA, spanB) => {
-          spanA.duration < spanB.duration;
-        });
-
-        const actualSpanNames: string[] = [];
-        spans.forEach(span => {
-          actualSpanNames.push(span.name);
-        });
-
-        const expectedSpanNames = [
-          'CloudSpanner.BatchTransaction.createPartitions_',
-          'CloudSpanner.BatchTransaction.createQueryPartitions',
-        ];
-        assert.deepStrictEqual(
-          actualSpanNames,
-          expectedSpanNames,
-          `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
-        );
-
-        // Ensure that createPartitions_ is a child span of createQueryPartitions.
-        const spanCreatePartitions_ = spans[0];
-        const spanCreateQueryPartitions = spans[1];
-        assert.ok(
-          spanCreateQueryPartitions.spanContext().traceId,
-          'Expected that createQueryPartitions has a defined traceId'
-        );
-        assert.ok(
-          spanCreatePartitions_.spanContext().traceId,
-          'Expected that createPartitions_ has a defined traceId'
-        );
-        assert.deepStrictEqual(
-          spanCreatePartitions_.spanContext().traceId,
-          spanCreateQueryPartitions.spanContext().traceId,
-          'Expected that both spans share a traceId'
-        );
-        assert.ok(
-          spanCreateQueryPartitions.spanContext().spanId,
-          'Expected that createQueryPartitions has a defined spanId'
-        );
-        assert.ok(
-          spanCreatePartitions_.spanContext().spanId,
-          'Expected that createPartitions_ has a defined spanId'
-        );
-        assert.deepStrictEqual(
-          spanCreatePartitions_.parentSpanId,
-          spanCreateQueryPartitions.spanContext().spanId,
-          'Expected that createQueryPartitions is the parent to createPartitions_'
-        );
-        done();
-      }
-    );
+    batchTransaction.createQueryPartitions(QUERY, err => {
+      assert.ifError(err);
+      traceExporter.forceFlush();
+      const spans = traceExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
+
+      // Sort the spans by duration.
+      spans.sort((spanA, spanB) => {
+        spanA.duration < spanB.duration;
+      });
+
+      const actualSpanNames: string[] = [];
+      spans.forEach(span => {
+        actualSpanNames.push(span.name);
+      });
+
+      const expectedSpanNames = [
+        'CloudSpanner.BatchTransaction.createPartitions_',
+        'CloudSpanner.BatchTransaction.createQueryPartitions',
+      ];
+      assert.deepStrictEqual(
+        actualSpanNames,
+        expectedSpanNames,
+        `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
+      );
+
+      // Ensure that createPartitions_ is a child span of createQueryPartitions.
+      const spanCreatePartitions_ = spans[0];
+      const spanCreateQueryPartitions = spans[1];
+      assert.ok(
+        spanCreateQueryPartitions.spanContext().traceId,
+        'Expected that createQueryPartitions has a defined traceId'
+      );
+      assert.ok(
+        spanCreatePartitions_.spanContext().traceId,
+        'Expected that createPartitions_ has a defined traceId'
+      );
+      assert.deepStrictEqual(
+        spanCreatePartitions_.spanContext().traceId,
+        spanCreateQueryPartitions.spanContext().traceId,
+        'Expected that both spans share a traceId'
+      );
+      assert.ok(
+        spanCreateQueryPartitions.spanContext().spanId,
+        'Expected that createQueryPartitions has a defined spanId'
+      );
+      assert.ok(
+        spanCreatePartitions_.spanContext().spanId,
+        'Expected that createPartitions_ has a defined spanId'
+      );
+      assert.deepStrictEqual(
+        spanCreatePartitions_.parentSpanId,
+        spanCreateQueryPartitions.spanContext().spanId,
+        'Expected that createQueryPartitions is the parent to createPartitions_'
+      );
+      done();
+    });
   });
 
   it('createReadPartitions', done => {
     const REQUEST = sandbox.stub();
     const response = {};
     REQUEST.callsFake((_, callback) => callback(null, response));
 
-    const res = batchTransaction.createReadPartitions(
-      QUERY,
-      (err, part, resp) => {
-        assert.ifError(err);
-        traceExporter.forceFlush();
-        const spans = traceExporter.getFinishedSpans();
-        assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
-
-        // Sort the spans by duration.
-        spans.sort((spanA, spanB) => {
-          spanA.duration < spanB.duration;
-        });
-
-        const actualSpanNames: string[] = [];
-        spans.forEach(span => {
-          actualSpanNames.push(span.name);
-        });
-        const expectedSpanNames = [
-          'CloudSpanner.BatchTransaction.createPartitions_',
-          'CloudSpanner.BatchTransaction.createReadPartitions',
-        ];
-        assert.deepStrictEqual(
-          actualSpanNames,
-          expectedSpanNames,
-          `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
-        );
-        done();
-      }
-    );
+    batchTransaction.createReadPartitions(QUERY, err => {
+      assert.ifError(err);
+      traceExporter.forceFlush();
+      const spans = traceExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
+
+      // Sort the spans by duration.
+      spans.sort((spanA, spanB) => {
+        spanA.duration < spanB.duration;
+      });
+
+      const actualSpanNames: string[] = [];
+      spans.forEach(span => {
+        actualSpanNames.push(span.name);
+      });
+      const expectedSpanNames = [
+        'CloudSpanner.BatchTransaction.createPartitions_',
+        'CloudSpanner.BatchTransaction.createReadPartitions',
+      ];
+      assert.deepStrictEqual(
+        actualSpanNames,
+        expectedSpanNames,
+        `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
+      );
+      done();
+    });
   });
 });
diff --git a/observability-test/database.ts b/observability-test/database.ts
@@ -507,7 +507,6 @@ describe('Database', () => {
 
     let beginSnapshotStub: sinon.SinonStub;
     let getSessionStub: sinon.SinonStub;
-    let snapshotStub: sinon.SinonStub;
 
     beforeEach(() => {
       fakePool = database.pool_;
@@ -524,9 +523,7 @@ describe('Database', () => {
         sandbox.stub(fakePool, 'getSession') as sinon.SinonStub
       ).callsFake(callback => callback(null, fakeSession));
 
-      snapshotStub = sandbox
-        .stub(fakeSession, 'snapshot')
-        .returns(fakeSnapshot);
+      sandbox.stub(fakeSession, 'snapshot').returns(fakeSnapshot);
     });
 
     it('with error', done => {
@@ -1175,7 +1172,7 @@ describe('Database', () => {
 
     it('with error on null mutation should catch thrown error', done => {
       try {
-        database.writeAtLeastOnce(null, (err, res) => {});
+        database.writeAtLeastOnce(null, () => {});
       } catch (err) {
         // Performing a substring search on the error because
         // depending on the version of Node.js, the error might be either of:
@@ -1250,7 +1247,6 @@ describe('Database', () => {
     let fakeSession: FakeSession;
     let fakeDataStream: Transform;
     let getSessionStub: sinon.SinonStub;
-    let requestStreamStub: sinon.SinonStub;
 
     const options = {
       requestOptions: {
@@ -1269,9 +1265,7 @@ describe('Database', () => {
         sandbox.stub(fakePool, 'getSession') as sinon.SinonStub
       ).callsFake(callback => callback(null, fakeSession));
 
-      requestStreamStub = sandbox
-        .stub(database, 'requestStream')
-        .returns(fakeDataStream);
+      sandbox.stub(database, 'requestStream').returns(fakeDataStream);
     });
 
     it('on retry with "Session not found" error', done => {
@@ -1320,7 +1314,6 @@ describe('Database', () => {
             'Expected an ERROR span status'
           );
 
-          const errorMessage = firstSpan.status.message;
           assert.deepStrictEqual(
             firstSpan.status.message,
             sessionNotFoundError.message
@@ -1658,7 +1651,7 @@ describe('Database', () => {
         .throws(ourException);
 
       assert.rejects(async () => {
-        const value = await database.runTransactionAsync(async txn => {
+        await database.runTransactionAsync(async txn => {
           const result = await txn.run('SELECT 1');
           await txn.commit();
           return result;
@@ -1724,8 +1717,6 @@ describe('Database', () => {
     let fakeStream2: Transform;
 
     let getSessionStub: sinon.SinonStub;
-    let snapshotStub: sinon.SinonStub;
-    let runStreamStub: sinon.SinonStub;
 
     beforeEach(() => {
       fakePool = database.pool_;
@@ -1746,15 +1737,11 @@ describe('Database', () => {
         .onSecondCall()
         .callsFake(callback => callback(null, fakeSession2));
 
-      snapshotStub = sandbox
-        .stub(fakeSession, 'snapshot')
-        .returns(fakeSnapshot);
+      sandbox.stub(fakeSession, 'snapshot').returns(fakeSnapshot);
 
       sandbox.stub(fakeSession2, 'snapshot').returns(fakeSnapshot2);
 
-      runStreamStub = sandbox
-        .stub(fakeSnapshot, 'runStream')
-        .returns(fakeStream);
+      sandbox.stub(fakeSnapshot, 'runStream').returns(fakeStream);
 
       sandbox.stub(fakeSnapshot2, 'runStream').returns(fakeStream2);
     });
@@ -1975,7 +1962,6 @@ describe('Database', () => {
 
     let getSessionStub;
     let beginStub;
-    let runUpdateStub;
 
     beforeEach(() => {
       fakePool = database.pool_;
@@ -1996,7 +1982,7 @@ describe('Database', () => {
         sandbox.stub(fakePartitionedDml, 'begin') as sinon.SinonStub
       ).callsFake(callback => callback(null));
 
-      runUpdateStub = (
+      (
         sandbox.stub(fakePartitionedDml, 'runUpdate') as sinon.SinonStub
       ).callsFake((_, callback) => callback(null));
     });
@@ -2031,7 +2017,6 @@ describe('Database', () => {
 
     it('with pool errors', done => {
       const fakeError = new Error('err');
-      const fakeCallback = sandbox.spy();
 
       getSessionStub.callsFake(callback => callback(fakeError));
       database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
@@ -2132,7 +2117,7 @@ describe('Database', () => {
         sandbox.stub(fakePool, 'release') as sinon.SinonStub
       ).withArgs(fakeSession);
 
-      database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
+      database.runPartitionedUpdate(QUERY, async () => {
         const exportResults = await getTraceExportResults();
         const actualSpanNames = exportResults.spanNames;
         const spans = exportResults.spans;

diff --git a/observability-test/helper.ts b/observability-test/helper.ts
@@ -19,6 +19,25 @@ import * as assert from 'assert';
 const {ReadableSpan} = require('@opentelemetry/sdk-trace-base');
 import {SEMATTRS_DB_NAME} from '@opentelemetry/semantic-conventions';
 
+export const batchCreateSessionsEvents = [
+  'Requesting 25 sessions',
+  'Creating 25 sessions',
+  'Requested for 25 sessions returned 25',
+];
+
+export const waitingSessionsEvents = [
+  'Acquiring session',
+  'Waiting for a session to become available',
+  'Acquired session',
+  'Using Session',
+];
+
+export const cacheSessionEvents = [
+  'Acquiring session',
+  'Cache hit: has usable session',
+  'Acquired session',
+];
+
 /**
  * This utility exists as a test helper because mocha has builtin "context"
  * and referring to context causes type/value collision errors.
@@ -47,3 +66,30 @@ export function generateWithAllSpansHaveDBName(dbName: String): Function {
     });
   };
 }
+
+export async function verifySpansAndEvents(
+  traceExporter,
+  expectedSpans,
+  expectedEvents
+) {
+  await traceExporter.forceFlush();
+  const spans = traceExporter.getFinishedSpans();
+  const actualEventNames: string[] = [];
+  const actualSpanNames: string[] = [];
+  spans.forEach(span => {
+    actualSpanNames.push(span.name);
+    span.events.forEach(event => {
+      actualEventNames.push(event.name);
+    });
+  });
+  assert.deepStrictEqual(
+    actualSpanNames,
+    expectedSpans,
+    `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpans}`
+  );
+  assert.deepStrictEqual(
+    actualEventNames,
+    expectedEvents,
+    `Unexpected events:\n\tGot:  ${actualEventNames}\n\tWant: ${expectedEvents}`
+  );
+}