From 01fdbc03519a518ca7047b7e1a9490e052afd92f Mon Sep 17 00:00:00 2001
From: Jason Myers <jason.myers@asurion.com>
Date: Thu, 11 Feb 2021 10:45:20 -0600
Subject: [PATCH] Adding support for serverless.[ts,json,js] files

---
 index.js | 47 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 8 deletions(-)

diff --git a/index.js b/index.js
index 29822bf..7bf0dc6 100644
--- a/index.js
+++ b/index.js
@@ -69,25 +69,56 @@ class ServerlessStepFunctionsLocal {
   }
 
   async getStepFunctionsFromConfig() {
-    const {servicePath} = this.serverless.config;
+    const fromYamlFile = (serverlessYmlPath) =>
+      this.serverless.yamlParser.parse(serverlessYmlPath);
 
-    if (!servicePath) {
-      throw new Error('service path not found');
-    }
+    let parsed = {};
+    let parser = null;
 
-    const configPath = path.join(servicePath, 'serverless.yml');
+    if (!this.serverless.service.stepFunctions) {
+      const { servicePath } = this.serverless.config;
 
-    const parsed = await this.serverless.yamlParser.parse(configPath);
+      if (!servicePath) {
+        throw new Error('service path not found');
+      }
+      const serviceFileName =
+        this.options.config ||
+        this.serverless.config.serverless.service.serviceFilename ||
+        'serverless.yml';
+      const configPath = path.join(servicePath, serviceFileName);
+      if (['.js', '.json', '.ts'].includes(path.extname(configPath))) {
+        parser = this.loadFromRequiredFile;
+      } else {
+        parser = fromYamlFile;
+      }
+      parsed = await parser(configPath);
+    } else {
+      parsed = this.serverless.service;
+    }
 
     this.stateMachines = parsed.stepFunctions.stateMachines;
 
     if (parsed.custom &&
       parsed.custom.stepFunctionsLocal &&
-      parsed.custom.stepFunctionsLocal.TaskResourceMapping) {
-      this.replaceTaskResourceMappings(parsed.stepFunctions.stateMachines, parsed.custom.stepFunctionsLocal.TaskResourceMapping);
+      parsed.custom.stepFunctionsLocal.TaskResourceMapping
+    ) {
+      this.replaceTaskResourceMappings(
+        parsed.stepFunctions.stateMachines,
+        parsed.custom.stepFunctionsLocal.TaskResourceMapping
+      );
     }
   }
 
+  // This function must be ignored since mocking the require system is more
+  // dangerous than beneficial
+  loadFromRequiredFile(serverlessYmlPath) {
+    /* istanbul ignore next */
+    // eslint-disable-next-line global-require, import/no-dynamic-require
+    const fileContents = require(serverlessYmlPath);
+    /* istanbul ignore next */
+    return Promise.resolve(fileContents);
+  }
+
   /**
    * Replaces Resource properties with values mapped in TaskResourceMapping
    */