diff --git a/src/assertions/toHaveCompletedExecutionWithStatus/index.ts b/src/assertions/toHaveCompletedExecutionWithStatus/index.ts
index a9c5075..4b683b7 100644
--- a/src/assertions/toHaveCompletedExecutionWithStatus/index.ts
+++ b/src/assertions/toHaveCompletedExecutionWithStatus/index.ts
@@ -1,4 +1,5 @@
 import { testResult, TestResultOutput } from "../../utils/testResult";
+import StepFunctions from "../../helpers/stepFunctions";
 
 import * as AWS from "aws-sdk";
 
@@ -8,16 +9,10 @@ export default {
     expectedStatus: string
   ): Promise<TestResultOutput> {
     const stepFunctions = new AWS.StepFunctions();
-    const listStateMachineParams = {};
-    // Get all state machines
-    const allStateMachines = await stepFunctions
-      .listStateMachines(listStateMachineParams)
-      .promise();
-    // Find state machine with specified name and get its arn
-    const smList = allStateMachines.stateMachines.filter(
-      (stateMachine: any) => stateMachine.name === stateMachineName
-    );
-    const smArn = smList[0].stateMachineArn;
+
+    const stepFunctionsObject =  await StepFunctions.build();
+    const smArn = await stepFunctionsObject.obtainStateMachineArn(stateMachineName); 
+    
     const listExecutionsParams = { stateMachineArn: smArn };
     // Get all executions of specified state machine
     const smExecutions = await stepFunctions
diff --git a/src/assertions/toMatchStateMachineOutput/index.ts b/src/assertions/toMatchStateMachineOutput/index.ts
index 56c4041..3b0e20f 100644
--- a/src/assertions/toMatchStateMachineOutput/index.ts
+++ b/src/assertions/toMatchStateMachineOutput/index.ts
@@ -1,5 +1,6 @@
-import { testResult, TestResultOutput } from "utils/testResult";
+import { testResult, TestResultOutput } from "../../utils/testResult";
 import { StepFunctions as AWSStepFunctions } from "aws-sdk";
+import StepFunctions from "../../helpers/stepFunctions";
 
 export default {
   async toMatchStateMachineOutput(
@@ -7,14 +8,9 @@ export default {
     expectedOutput: any
   ): Promise<TestResultOutput> {
     const stepFunctions = new AWSStepFunctions();
-    const allStateMachines = await stepFunctions.listStateMachines().promise();
-
-    const smList = allStateMachines.stateMachines.filter(
-      (stateMachine: any) => stateMachine.name === stateMachineName
-    );
-
-    const stateMachineArn = smList[0].stateMachineArn;
-    const listExecParams = { stateMachineArn: stateMachineArn };
+    const stepFunctionsObject =  await StepFunctions.build();
+    const smArn = await stepFunctionsObject.obtainStateMachineArn(stateMachineName); 
+    const listExecParams = { stateMachineArn: smArn };
     const executionList = await stepFunctions
       .listExecutions(listExecParams)
       .promise();
diff --git a/src/helpers/stepFunctions.ts b/src/helpers/stepFunctions.ts
index 5e31d35..118bf8d 100644
--- a/src/helpers/stepFunctions.ts
+++ b/src/helpers/stepFunctions.ts
@@ -61,4 +61,27 @@ export default class StepFunctions {
       .describeExecution({ executionArn: execution.executionArn })
       .promise();
   }
+
+  async obtainStateMachineArn(stateMachineName: string): Promise<string> {
+    
+    const listStateMachineParams = {};
+    // Get all state machines
+    if (this.stepFunctions === undefined) {
+      throw new Error(
+        "The Step Functions client is undefined. You might have forgotten to run build()."
+      );
+    }
+    const allStateMachines = await this.stepFunctions
+      .listStateMachines(listStateMachineParams)
+      .promise();
+    // Find state machine with specified name and get its arn
+    const smList = allStateMachines.stateMachines.find(
+      (stateMachine: any) => stateMachine.name === stateMachineName
+    );
+    if (smList == null)
+    throw new Error(
+      "No matching state machine. "
+    );
+    return smList.stateMachineArn;
+  }
 }