diff --git a/packages/server/src/sdk/app/automations/tests/index.spec.ts b/packages/server/src/sdk/app/automations/tests/index.spec.ts
index 295ab690d1..124a5d9276 100644
--- a/packages/server/src/sdk/app/automations/tests/index.spec.ts
+++ b/packages/server/src/sdk/app/automations/tests/index.spec.ts
@@ -1,4 +1,4 @@
-import _ from "lodash/fp"
+import { sample } from "lodash/fp"
 import { Automation } from "@budibase/types"
 import automationSdk from "../"
 import { structures } from "../../../../api/routes/tests/utilities"
@@ -12,39 +12,40 @@ describe("automation sdk", () => {
   })
 
   describe("update", () => {
-    it("can update input fields", async () => {
+    it.each([
+      ["trigger", (a: Automation) => a.definition.trigger],
+      ["step", (a: Automation) => a.definition.steps[0]],
+    ])("can update input fields (for a %s)", async (_, getStep) => {
       await config.doInContext(config.getAppId(), async () => {
-        const automation: Automation = structures.newAutomation()
-        const keyToUse = _.sample(
-          Object.keys(automation.definition.trigger.inputs)
-        )!
-        automation.definition.trigger.inputs[keyToUse] = "anyValue"
+        const automation = structures.newAutomation()
+
+        const keyToUse = sample(Object.keys(getStep(automation).inputs))!
+        getStep(automation).inputs[keyToUse] = "anyValue"
 
         const response = await automationSdk.create(automation)
 
         const update = { ...response }
-        update.definition.trigger.inputs[keyToUse] = "anyUpdatedValue"
+        getStep(update).inputs[keyToUse] = "anyUpdatedValue"
         const result = await automationSdk.update(update)
-        expect(result.definition.trigger.inputs[keyToUse]).toEqual(
-          "anyUpdatedValue"
-        )
+        expect(getStep(result).inputs[keyToUse]).toEqual("anyUpdatedValue")
       })
     })
 
-    it("cannot update readonly fields", async () => {
+    it.each([
+      ["trigger", (a: Automation) => a.definition.trigger],
+      ["step", (a: Automation) => a.definition.steps[0]],
+    ])("cannot update readonly fields (for a %s)", async (_, getStep) => {
       await config.doInContext(config.getAppId(), async () => {
-        const automation: Automation = { ...structures.newAutomation() }
-        automation.definition.trigger.schema.inputs.properties[
-          "readonlyProperty"
-        ] = {
+        const automation = structures.newAutomation()
+        getStep(automation).schema.inputs.properties["readonlyProperty"] = {
           readonly: true,
         }
-        automation.definition.trigger.inputs["readonlyProperty"] = "anyValue"
+        getStep(automation).inputs["readonlyProperty"] = "anyValue"
 
         const response = await automationSdk.create(automation)
 
         const update = { ...response }
-        update.definition.trigger.inputs["readonlyProperty"] = "anyUpdatedValue"
+        getStep(update).inputs["readonlyProperty"] = "anyUpdatedValue"
         await expect(automationSdk.update(update)).rejects.toThrow(
           "Field readonlyProperty is readonly and it cannot be modified"
         )