[pinpoint-apm#251] GrpcDataSender fixtures close resources

feelform · Dec 23, 2024 · 9d680c3 · 9d680c3
1 parent e13ef31
commit 9d680c3
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 33 deletions.
diff --git a/test/client/mock-grpc-data-sender.js b/test/client/mock-grpc-data-sender.js
@@ -13,6 +13,8 @@ class MockGrpcDataSender extends GrpcDataSender {
     this.agentClient = {
       requestAgentInfo: function (pAgentInfo) {
         self.actualAgentInfo = pAgentInfo
+      },
+      close: function () {
       }
     }
 
@@ -28,12 +30,14 @@ class MockGrpcDataSender extends GrpcDataSender {
       },
       requestSqlUidMetaData: function (pSqlUidMetaData) {
         self.actualSqlUidMetaData = pSqlUidMetaData
+      },
+      close: function () {
       }
     }
     this.actualSpans = []
   }
 
-  get actualSpan () {
+  get actualSpan() {
     return this.actualSpans[this.actualSpans.length - 1]
   }
 
@@ -43,25 +47,32 @@ class MockGrpcDataSender extends GrpcDataSender {
   initializeSpanStream() {
     let self = this
     this.spanStream = {
-      write: function (span) {
-        self.actualSpans.push(span)
-      },
-      end: function () {
+      grpcStream: {
+        write: function (span) {
+          self.actualSpans.push(span)
+        },
+        end: function () {
 
+        }
       }
     }
   }
 
   initializeProfilerClients() {
     let self = this
     this.commandStream = {
-      write: function (pmessage) {
-        self.actualPCmdMessage = pmessage
-      },
-      end: function () {
+      grpcStream: {
+        write: function (pmessage) {
+          self.actualPCmdMessage = pmessage
+        },
+        end: function () {
+
+        },
+        on: function () {
 
+        }
       },
-      on: function () {
+      end() {
 
       }
     }
@@ -70,29 +81,33 @@ class MockGrpcDataSender extends GrpcDataSender {
   initializeStatStream() {
     let self = this
     this.statStream = {
-      write: function (pmessage) {
-        self.actualPStatMessage = pmessage
-      },
-      end: function () {
+      grpcStream: {
+        write: function (pmessage) {
+          self.actualPStatMessage = pmessage
+        },
+        end: function () {
 
-      },
-      on: function () {
+        },
+        on: function () {
 
+        }
       }
     }
   }
 
   initializePingStream() {
     let self = this
     this.pingStream = {
-      write: function (pmessage) {
-        self.actualPingMessage = pmessage
-      },
-      end: function () {
+      grpcStream: {
+        write: function (pmessage) {
+          self.actualPingMessage = pmessage
+        },
+        end: function () {
 
-      },
-      on: function () {
+        },
+        on: function () {
 
+        }
       }
     }
   }

diff --git a/test/instrumentation/context/nested-async-trace.test.js b/test/instrumentation/context/nested-async-trace.test.js
@@ -57,8 +57,7 @@ const spanMessageStreamService = (call) => {
         }
     })
 
-    call.on('end', () => {
-    })
+    call.on('end', () => {})
 }
 
 function sendAgentStat(call, callback) {
@@ -95,6 +94,7 @@ test(`nested mysql async query with express`, async (t) => {
     collectorServer.addService(services.MetadataService, {
         requestApiMetaData: pResultReturnUnaryService,
         requestSqlUidMetaData: pResultReturnUnaryService,
+        requestSqlMetaData: pResultReturnUnaryService
     })
     process.env['PINPOINT_PROFILER_SQL_STAT'] = 'true'
 

diff --git a/test/support/agent-singleton-mock.js b/test/support/agent-singleton-mock.js
@@ -109,10 +109,12 @@ class MockAgent extends Agent {
 
         const dataSender = this.makeDataSender()
         this.traceContext.dataSender = dataSender
+        this.dataSender.close()
         this.dataSender = dataSender
         // this.initializeDataSender(dataSender)
         stringMetaService.init(dataSender)
         apiMetaService.init(dataSender)
+        sqlMetaDataService.setDataSender(dataSender)
 
         resetSpanOrSpanChunks()
 

diff --git a/test/support/data-sender-mock.js b/test/support/data-sender-mock.js
@@ -15,13 +15,16 @@ const DataSender = require('../../lib/client/data-sender')
 const MockGrpcDataSender = require('../client/mock-grpc-data-sender')
 const SqlMetaData = require('../../lib/client/sql-meta-data')
 const GrpcDataSender = require('../../lib/client/grpc-data-sender')
+const SqlUidMetaData = require('../../lib/client/sql-uid-meta-data')
 
 class MockDataSender extends DataSender {
   constructor(config, dataSender) {
     super(config, dataSender)
     this.mockAPIMetaInfos = []
     this.mockSpanChunks = []
     this.mockSpans = []
+    this.mockSqlMetadata = []
+    this.mockSqlUidMetadata = []
   }
 
   send(data) {
@@ -31,31 +34,26 @@ class MockDataSender extends DataSender {
 
     if (data instanceof AgentInfo) {
       this.mockAgentInfo = data
-      super.send(data)
     } else if (data instanceof ApiMetaInfo) {
       this.mockAPIMetaInfos.push(data)
-      super.send(data)
     } else if (data instanceof StringMetaInfo) {
       this.mockMetaInfo = data
-      super.send(data)
     } else if (data instanceof Span) {
       this.mockSpan = data
       this.mockSpans.push(data)
-      super.send(data)
     } else if (data instanceof SpanChunk) {
       this.mockSpanChunks.push(data)
-      super.send(data)
     } else if (data instanceof SqlMetaData) {
-      this.mockSqlMetaData = data
-      super.send(data)
+      this.mockSqlMetadata.push(data)
     } else if (data?.isAsyncSpanChunk?.()) {
       this.mockSpanChunks.push(data)
-      super.send(data)
     } else if (data?.isSpan?.()) {
       this.mockSpan = data
       this.mockSpans.push(data)
-      super.send(data)
+    } else if (data instanceof SqlUidMetaData) {
+      this.mockSqlUidMetadata.push(data)
     }
+    super.send(data)
   }
 
   findSpanChunk(asyncId) {
@@ -69,6 +67,8 @@ class MockDataSender extends DataSender {
     this.mockAPIMetaInfos = []
     this.mockSpanChunks = []
     this.mockSpans = []
+    this.mockSqlMetadata = []
+    this.mockSqlUidMetadata = []
   }
 
   close() {