From 31e5d076d3938ad6dcbd5968f11d0cdb986a2620 Mon Sep 17 00:00:00 2001
From: horike37 <horike37@gmail.com>
Date: Sat, 18 Feb 2017 20:01:59 +0900
Subject: [PATCH] replacing task arn

---
 lib/dataProcessing.js      | 32 ++++++++++++++----
 lib/dataProcessing.test.js | 69 ++++++++++++++++++++++++++++++++++++++
 lib/index.js               |  1 +
 lib/index.test.js          |  6 +++-
 4 files changed, 101 insertions(+), 7 deletions(-)

diff --git a/lib/dataProcessing.js b/lib/dataProcessing.js
index 92509964..22dd8228 100644
--- a/lib/dataProcessing.js
+++ b/lib/dataProcessing.js
@@ -72,11 +72,21 @@ module.exports = {
 
     this.serverless.service.stepFunctions.stateMachines[this.options.name] =
       JSON.stringify(this.serverless.service.stepFunctions.stateMachines[this.options.name]);
-    _.forEach(this.functionArns, (value, key) => {
-      const regExp = new RegExp(`"Resource":"${key}"`, 'g');
+    _.forEach(this.functionArns, (functionArn, functionName) => {
+      const regExp = new RegExp(`"Resource":"${functionName}"`, 'g');
       this.serverless.service.stepFunctions.stateMachines[this.options.name] =
         this.serverless.service.stepFunctions.stateMachines[this.options.name]
-        .replace(regExp, `"Resource":"${value}"`);
+        .replace(regExp, `"Resource":"${functionArn}"`);
+    });
+
+    _.forEach(this.activityArns, (activityArn, activityName) => {
+      const regExp = new RegExp(`"Resource":"${activityName}"`, 'g');
+      _.forEach(this.serverless.service.stepFunctions.stateMachines,
+        (stepFunctionObj, stepFunctionKey) => {
+          this.serverless.service.stepFunctions.stateMachines[stepFunctionKey] =
+          this.serverless.service.stepFunctions.stateMachines[stepFunctionKey]
+          .replace(regExp, `"Resource":"${activityArn}"`);
+        });
     });
     return BbPromise.resolve();
   },
@@ -95,13 +105,23 @@ module.exports = {
         = JSON.stringify(stepFunctionObj);
       });
 
-    _.forEach(this.functionArns, (functionObj, functionKey) => {
-      const regExp = new RegExp(`"Resource":"${functionKey}"`, 'g');
+    _.forEach(this.functionArns, (functionArn, functionName) => {
+      const regExp = new RegExp(`"Resource":"${functionName}"`, 'g');
+      _.forEach(this.serverless.service.stepFunctions.stateMachines,
+        (stepFunctionObj, stepFunctionKey) => {
+          this.serverless.service.stepFunctions.stateMachines[stepFunctionKey] =
+          this.serverless.service.stepFunctions.stateMachines[stepFunctionKey]
+          .replace(regExp, `"Resource":"${functionArn}"`);
+        });
+    });
+
+    _.forEach(this.activityArns, (activityArn, activityName) => {
+      const regExp = new RegExp(`"Resource":"${activityName}"`, 'g');
       _.forEach(this.serverless.service.stepFunctions.stateMachines,
         (stepFunctionObj, stepFunctionKey) => {
           this.serverless.service.stepFunctions.stateMachines[stepFunctionKey] =
           this.serverless.service.stepFunctions.stateMachines[stepFunctionKey]
-          .replace(regExp, `"Resource":"${functionObj}"`);
+          .replace(regExp, `"Resource":"${activityArn}"`);
         });
     });
     return BbPromise.resolve();
diff --git a/lib/dataProcessing.test.js b/lib/dataProcessing.test.js
index 6505f212..e83296c6 100644
--- a/lib/dataProcessing.test.js
+++ b/lib/dataProcessing.test.js
@@ -142,6 +142,49 @@ describe('dataProsessing', () => {
         .to.be.equal(a);
       });
     });
+
+    it('should comple with replacing activityArn', () => {
+      serverless.service.stepFunctions.stateMachines = {
+        hellofunc: {
+          States: {
+            HelloWorld: {
+              Resource: 'someActivity',
+            },
+          },
+        },
+      };
+      serverlessStepFunctions.activityArns.someActivity = 'activityArn';
+      serverlessStepFunctions.compile().then(() => {
+        expect(serverlessStepFunctions.serverless.service.stepFunctions.stateMachines.hellofunc)
+        .to.be.equal('{"States":{"HelloWorld":{"Resource":"activityArn"}}}');
+      });
+    });
+
+    it('should comple with correct params when nested Resource with replacing activityArn', () => {
+      serverlessStepFunctions.serverless.service.stepFunctions.stateMachines = {
+        hellofunc: {
+          States: {
+            HelloWorld: {
+              Resource: 'someActivity',
+              HelloWorld: {
+                Resource: 'someActivity',
+                HelloWorld: {
+                  Resource: 'someActivity',
+                },
+              },
+            },
+          },
+        },
+      };
+
+      let a = '{"States":{"HelloWorld":{"Resource":"activityArn","HelloWorld"';
+      a += ':{"Resource":"activityArn","HelloWorld":{"Resource":"activityArn"}}}}}';
+      serverlessStepFunctions.activityArns.someActivity = 'activityArn';
+      serverlessStepFunctions.compile().then(() => {
+        expect(serverlessStepFunctions.serverless.service.stepFunctions.stateMachines.hellofunc)
+        .to.be.equal(a);
+      });
+    });
   });
 
   describe('#parseInputdate()', () => {
@@ -242,5 +285,31 @@ describe('dataProsessing', () => {
         .to.be.equal(a);
       });
     });
+
+    it('should comple with correct params when nested Resource with replacing activityArn', () => {
+      serverlessStepFunctions.serverless.service.stepFunctions.stateMachines = {
+        hellofunc: {
+          States: {
+            HelloWorld: {
+              Resource: 'someActivity',
+              HelloWorld: {
+                Resource: 'someActivity',
+                HelloWorld: {
+                  Resource: 'someActivity',
+                },
+              },
+            },
+          },
+        },
+      };
+
+      let a = '{"States":{"HelloWorld":{"Resource":"activityArn","HelloWorld"';
+      a += ':{"Resource":"activityArn","HelloWorld":{"Resource":"activityArn"}}}}}';
+      serverlessStepFunctions.activityArns.someActivity = 'activityArn';
+      serverlessStepFunctions.compileAll().then(() => {
+        expect(serverlessStepFunctions.serverless.service.stepFunctions.stateMachines.hellofunc)
+        .to.be.equal(a);
+      });
+    });
   });
 });
diff --git a/lib/index.js b/lib/index.js
index 9828f5ba..28c56b96 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -354,6 +354,7 @@ class ServerlessStepFunctions {
       .then(this.yamlParse)
       .then(this.getStateMachineArn)
       .then(this.getFunctionArns)
+      .then(this.getActivityArns)
       .then(this.compile)
       .then(this.getIamRole)
       .then(this.deleteStateMachine)
diff --git a/lib/index.test.js b/lib/index.test.js
index 837182d4..31c4c88e 100644
--- a/lib/index.test.js
+++ b/lib/index.test.js
@@ -252,6 +252,8 @@ describe('ServerlessStepFunctions', () => {
         .stub(serverlessStepFunctions, 'getStateMachineArn').returns(BbPromise.resolve());
       const getFunctionArnsStub = sinon
         .stub(serverlessStepFunctions, 'getFunctionArns').returns(BbPromise.resolve());
+      const getActivityArnsStub = sinon
+        .stub(serverlessStepFunctions, 'getActivityArns').returns(BbPromise.resolve());
       const compileStub = sinon
         .stub(serverlessStepFunctions, 'compile').returns(BbPromise.resolve());
       const getIamRoleStub = sinon
@@ -266,7 +268,8 @@ describe('ServerlessStepFunctions', () => {
           expect(yamlParseStub.calledOnce).to.be.equal(true);
           expect(getStateMachineArnStub.calledAfter(yamlParseStub)).to.be.equal(true);
           expect(getFunctionArnsStub.calledAfter(getStateMachineArnStub)).to.be.equal(true);
-          expect(compileStub.calledAfter(getFunctionArnsStub)).to.be.equal(true);
+          expect(getActivityArnsStub.calledAfter(getFunctionArnsStub)).to.be.equal(true);
+          expect(compileStub.calledAfter(getActivityArnsStub)).to.be.equal(true);
           expect(getIamRoleStub.calledAfter(compileStub)).to.be.equal(true);
           expect(deleteStateMachineStub.calledAfter(getIamRoleStub)).to.be.equal(true);
           expect(createStateMachineStub.calledAfter(deleteStateMachineStub)).to.be.equal(true);
@@ -274,6 +277,7 @@ describe('ServerlessStepFunctions', () => {
           serverlessStepFunctions.yamlParse.restore();
           serverlessStepFunctions.getStateMachineArn.restore();
           serverlessStepFunctions.getFunctionArns.restore();
+          serverlessStepFunctions.getActivityArns.restore();
           serverlessStepFunctions.compile.restore();
           serverlessStepFunctions.getIamRole.restore();
           serverlessStepFunctions.deleteStateMachine.restore();