[Step Function] 3.3. Add class StepFunctionConfigLoader

serverless/src/index.ts

@@ -2,6 +2,7 @@ import { validateParameters as validateLambdaParameters, LambdaConfigLoader } fr
 import { instrumentLambdas } from "./lambda/lambda";
 import { InputEvent, OutputEvent, SUCCESS, FAILURE } from "./types";
 import { instrumentStateMachines } from "./step_function/step_function";
+import { StepFunctionConfigLoader } from "./step_function/env";
 import log from "loglevel";
 
 export const handler = async (event: InputEvent, _: any): Promise<OutputEvent> => {
@@ -27,7 +28,9 @@ export const handler = async (event: InputEvent, _: any): Promise<OutputEvent> =
       return lambdaOutput;
     }
 
-    const stepFunctionOutput = await instrumentStateMachines(event);
+    const stepFunctionConfig = new StepFunctionConfigLoader().getConfig(event);
+
+    const stepFunctionOutput = await instrumentStateMachines(event, stepFunctionConfig);
     return stepFunctionOutput;
   } catch (error: any) {
     return {

serverless/src/step_function/env.ts

@@ -0,0 +1,24 @@
+import { ConfigLoader } from "../env";
+
+export interface Configuration {
+  // When set, it will be added to the state machine's log group name.
+  env?: string;
+}
+
+const envEnvVar = "DD_ENV";
+
+export class StepFunctionConfigLoader extends ConfigLoader<Configuration> {
+  readonly defaultConfiguration: Configuration = {};
+
+  public getConfigFromEnvVars(): Configuration {
+    const config: Configuration = {
+      ...this.defaultConfiguration,
+    };
+
+    if (envEnvVar in process.env) {
+      config.env = process.env[envEnvVar];
+    }
+
+    return config;
+  }
+}

serverless/src/step_function/step_function.ts

@@ -1,17 +1,18 @@
 import { InputEvent, OutputEvent, SUCCESS, Resources } from "../types";
 import log from "loglevel";
 import { StateMachine, StateMachineProperties } from "./types";
+import { Configuration } from "./env";
 import { setUpLogging } from "./log";
 
 const STATE_MACHINE_RESOURCE_TYPE = "AWS::StepFunctions::StateMachine";
 
-export async function instrumentStateMachines(event: InputEvent): Promise<OutputEvent> {
+export async function instrumentStateMachines(event: InputEvent, config: Configuration): Promise<OutputEvent> {
   const fragment = event.fragment;
   const resources = fragment.Resources;
 
   const stateMachines = findStateMachines(resources);
   for (const stateMachine of stateMachines) {
-    instrumentStateMachine(resources, stateMachine);
+    instrumentStateMachine(resources, config, stateMachine);
   }
 
   return {
@@ -21,7 +22,7 @@ export async function instrumentStateMachines(event: InputEvent): Promise<Output
   };
 }
 
-function instrumentStateMachine(resources: Resources, stateMachine: StateMachine): void {
+function instrumentStateMachine(resources: Resources, config: Configuration, stateMachine: StateMachine): void {
   log.debug(`Instrumenting State Machine ${stateMachine.resourceKey}`);
 
   setUpLogging(resources, stateMachine);

serverless/test/step_function/env.spec.ts

@@ -0,0 +1,86 @@
+import { InputEvent } from "../../src/types";
+import { StepFunctionConfigLoader } from "../../src/step_function/env";
+
+const loader = new StepFunctionConfigLoader();
+describe("getConfig", () => {
+  const originalEnv = process.env;
+
+  beforeEach(async () => {
+    jest.resetModules();
+    process.env = { ...originalEnv };
+  });
+
+  afterEach(() => {
+    jest.resetAllMocks();
+    process.env = originalEnv;
+  });
+
+  // TODO: Change "env" to a field which is in defaultConfig once there is one
+  describe("1. CloudFormation Macro params are set", () => {
+    it("uses CloudFormation Macro params over environment variables and default configuration", () => {
+      const event: InputEvent = {
+        params: {
+          env: "macroEnv",
+        },
+        fragment: {
+          Mappings: {},
+        },
+      } as any;
+
+      process.env.DD_ENV = "envVarEnv";
+
+      const config = loader.getConfig(event);
+      expect(config.env).toBe("macroEnv");
+    });
+  });
+
+  describe("2. CloudFormation Mappings are set", () => {
+    it("uses CloudFormation Mappings params over environment variables and default configuration", () => {
+      const event: InputEvent = {
+        params: {},
+        fragment: {
+          Mappings: {
+            Datadog: {
+              Parameters: {
+                env: "mappingEnv",
+              },
+            },
+          },
+        },
+      } as any;
+
+      process.env.DD_ENV = "envVarEnv";
+
+      const config = loader.getConfig(event);
+      expect(config.env).toBe("mappingEnv");
+    });
+  });
+
+  describe("3. Neither CloudFormation Macro params nor CloudFormation Mappings is set", () => {
+    it("uses environment variables over default configuration", () => {
+      const event: InputEvent = {
+        params: {},
+        fragment: {
+          Mappings: {},
+        },
+      } as any;
+
+      process.env.DD_ENV = "envVarEnv";
+
+      const config = loader.getConfig(event);
+      expect(config.env).toBe("envVarEnv");
+    });
+
+    it("uses default configuration if no other params are set", () => {
+      const event: InputEvent = {
+        params: {},
+        fragment: {
+          Mappings: {},
+        },
+      } as any;
+
+      const config = loader.getConfig(event);
+      expect(config.env).toBe(undefined);
+    });
+  });
+});