diff --git a/test/client/mock-grpc-data-sender.js b/test/client/mock-grpc-data-sender.js index 2c223af1..bbfe1fce 100644 --- a/test/client/mock-grpc-data-sender.js +++ b/test/client/mock-grpc-data-sender.js @@ -7,6 +7,19 @@ 'use strict' const GrpcDataSender = require('../../lib/client/grpc-data-sender') +class MockGrpcStream { + constructor(stream) { + this.grpcStream = stream + } + + write(data) { + this.grpcStream.write(data) + } + + end() { + } +} + class MockGrpcDataSender extends GrpcDataSender { initializeClients() { let self = this @@ -46,70 +59,59 @@ class MockGrpcDataSender extends GrpcDataSender { initializeSpanStream() { let self = this - this.spanStream = { - grpcStream: { - write: function (span) { - self.actualSpans.push(span) - }, - end: function () { - - } + this.spanStream = new MockGrpcStream({ + write: function (span) { + self.actualSpans.push(span) + }, + end: function () { + } - } + }) } initializeProfilerClients() { let self = this - this.commandStream = { - grpcStream: { - write: function (pmessage) { - self.actualPCmdMessage = pmessage - }, - end: function () { - - }, - on: function () { + this.commandStream = new MockGrpcStream({ + write: function (pmessage) { + self.actualPCmdMessage = pmessage + }, + end: function () { - } }, - end() { + on: function () { } - } + }) } initializeStatStream() { let self = this - this.statStream = { - grpcStream: { - write: function (pmessage) { - self.actualPStatMessage = pmessage - }, - end: function () { + this.statStream = new MockGrpcStream({ + write: function (pmessage) { + self.actualPStatMessage = pmessage + }, + end: function () { - }, - on: function () { + }, + on: function () { - } } - } + }) } initializePingStream() { let self = this - this.pingStream = { - grpcStream: { - write: function (pmessage) { - self.actualPingMessage = pmessage - }, - end: function () { + this.pingStream = new MockGrpcStream({ + write: function (pmessage) { + self.actualPingMessage = pmessage + }, + end: function () { - }, - on: function () { + }, + on: function () { - } } - } + }) } initializeAgentInfoScheduler() {