Wrap the sendColector status function to send the status to DD metrics

rakhimundhada15 · Jul 10, 2023 · b78239c · b78239c
1 parent fb7055e
commit b78239c
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 0 deletions.
diff --git a/paws_collector.js b/paws_collector.js
@@ -224,6 +224,22 @@ class PawsCollector extends AlAwsCollector {
     prepareErrorStatus(errorString, collectorType = this.pawsCollectorType) {
         return super.setCollectorStatus(collectorType, errorString);
     }
+    /**
+     * Override the super method to send the status to CloudWatch and DataDog metrics
+     * @param {*} stream 
+     * @param {*} collectorStatus 
+     * @param {*} callback 
+     */
+    sendCollectorStatus(stream, collectorStatus, callback) {
+        super.sendCollectorStatus(stream, collectorStatus, (err) => {
+            // report status to metrics if collector_status send the responce succesful.
+            if (!err) {
+                return this.reportCollectorStatus(collectorStatus.status, callback);
+            }
+            else return callback(err);
+        }
+        );
+    }
 
     setPawsSecret(secretValue){
         const encryptPromise = new Promise((resolve, reject) => {
@@ -739,6 +755,40 @@ class PawsCollector extends AlAwsCollector {
         this.reportDDMetric("duplicate_messages", duplicateCount);
         return cloudwatch.putMetricData(params, callback);
     };
+    /**
+     * Report the collector status(ok/error)to dd metrics
+     * @param {*} status 
+     * @param {*} callback 
+     * @returns 
+     */
+    reportCollectorStatus(status, callback) {
+        var cloudwatch = new AWS.CloudWatch({apiVersion: '2010-08-01'});
+        const params = {
+            MetricData: [
+              {
+                MetricName: 'PawsCollectorStatus',
+                Dimensions: [
+                  {
+                    Name: 'CollectorType',
+                    Value: this._pawsCollectorType
+                  },
+                  {
+                    Name: 'FunctionName',
+                    Value: process.env.AWS_LAMBDA_FUNCTION_NAME
+                  }
+                ],
+                Timestamp: new Date(),
+                Unit: 'Count',
+                Value: 1
+              }
+            ],
+            Namespace: 'PawsCollectors'
+        };
+
+        this.reportDDMetric("collector_status", 1, [`status:${status}`]);
+        return cloudwatch.putMetricData(params, callback);
+    };
+
     _storeCollectionState(pawsState, privCollectorState, invocationTimeout, callback) {
         if (Array.isArray(privCollectorState)) {
             return this._storeCollectionStateArray(pawsState, privCollectorState, invocationTimeout, callback);

diff --git a/test/paws_test.js b/test/paws_test.js
@@ -116,6 +116,10 @@ function mockDDB(getItemStub, putItemStub, updateItemStub, batchWriteItemStub) {
     AWS.mock('DynamoDB', 'batchWriteItem', batchWriteItemStub ? batchWriteItemStub : defaultMock);
 }
 
+function mockCloudWatch() {
+    AWS.mock('CloudWatch', 'putMetricData', (params, callback) => callback());
+}
+
 function gen_state_objects(num) {
     return new Array(num).fill(0).map((e,i) => ({state: 'new-state-' + i}));
 }
@@ -705,6 +709,7 @@ describe('Unit Tests', function() {
 
         it('Check sendCollectorStatus method call only after Five failed attempt', function (done) {
             mockDDB();
+            mockCloudWatch();
             let ctx = {
                 invokedFunctionArn: pawsMock.FUNCTION_ARN,
                 fail: function (error) {
@@ -716,6 +721,7 @@ describe('Unit Tests', function() {
                     sinon.assert.calledOnce(mockPawsGetLogs);
                     mockPawsGetLogs.restore();
                     mockSendCollectorStatus.restore();
+                    AWS.restore('CloudWatch');
                     done();
                 }
             };
@@ -746,6 +752,7 @@ describe('Unit Tests', function() {
 
         it('Check sendCollectorStatus method not call if failed attempt less < 5', function (done) {
             mockDDB();
+            mockCloudWatch();
             let ctx = {
                 invokedFunctionArn: pawsMock.FUNCTION_ARN,
                 fail: function (error) {
@@ -757,6 +764,7 @@ describe('Unit Tests', function() {
                     sinon.assert.calledOnce(mockPawsGetLogs);
                     mockPawsGetLogs.restore();
                     mockSendCollectorStatus.restore();
+                    AWS.restore('CloudWatch');
                     done();
                 }
             };
@@ -787,6 +795,7 @@ describe('Unit Tests', function() {
 
         it('Check if retry_count get added in state for existing collector and it will not call mockSendCollectorStatus method', function (done) {
             mockDDB();
+            mockCloudWatch();
             let ctx = {
                 invokedFunctionArn: pawsMock.FUNCTION_ARN,
                 fail: function (error) {
@@ -798,6 +807,7 @@ describe('Unit Tests', function() {
                     sinon.assert.calledOnce(mockPawsGetLogs);
                     mockPawsGetLogs.restore();
                     mockSendCollectorStatus.restore();
+                    AWS.restore('CloudWatch');
                     done();
                 }
             };
@@ -1066,6 +1076,30 @@ describe('Unit Tests', function() {
                 });
             });
         });
+
+        it('reportCollectorStatus', function(done) {
+            let ctx = {
+                invokedFunctionArn : pawsMock.FUNCTION_ARN,
+                fail : function(error) {
+                    assert.fail(error);
+                    done();
+                },
+                succeed : function() {
+                    done();
+                }
+            };
+            AWS.mock('CloudWatch', 'putMetricData', (params, callback) => callback());
+            TestCollector.load().then(function(creds) {
+                var collector = new TestCollector(ctx, creds);
+                const status = 'ok';
+                collector.reportCollectorStatus(status, function(error) {
+                    assert.equal(null, error);
+                    AWS.restore('KMS');
+                    AWS.restore('CloudWatch');
+                    done();
+                });
+            });
+        });
     });
 
     describe('Register Tests', function() {