diff --git a/src/CqlBuilderPanel/functionsSection/argumentSection/ArgumentSection.tsx b/src/CqlBuilderPanel/functionsSection/argumentSection/ArgumentSection.tsx
index 04bedf4..f084adf 100644
--- a/src/CqlBuilderPanel/functionsSection/argumentSection/ArgumentSection.tsx
+++ b/src/CqlBuilderPanel/functionsSection/argumentSection/ArgumentSection.tsx
@@ -95,11 +95,11 @@ export default function ArgumentSection(props: ArgumentsProps) {
         <div tw="flex-grow pl-5">
           <Select
             label="Available DataTypes"
-            id="type-selector"
+            id="arg-type-selector"
             inputProps={{
-              "data-testid": "type-selector-input",
+              "data-testid": "arg-type-selector-input",
             }}
-            data-testid="type-selector"
+            data-testid="arg-type-selector"
             SelectDisplayProps={{
               "aria-required": "true",
             }}
diff --git a/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.test.tsx b/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.test.tsx
index 49ad14d..be73aa5 100644
--- a/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.test.tsx
+++ b/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.test.tsx
@@ -428,6 +428,136 @@ describe("CQL Function Builder Tests", () => {
     )) as HTMLInputElement;
     expect(definitionName.value).toBe("IP");
   });
+  it("should call handleApplyFunction with a function that we've created through the UI", async () => {
+    const handleApplyFn = jest.fn();
+    render(
+      <FunctionBuilder
+        canEdit={true}
+        handleApplyFunction={handleApplyFn}
+        cqlBuilderLookupsTypes={cqlBuilderLookup}
+      />
+    );
+    // name, insert, except args
+    const functionNameInput = (await screen.findByTestId(
+      "function-name-text-input"
+    )) as HTMLInputElement;
+    const argumentsSection = screen.getByTestId(
+      "terminology-section-Arguments-sub-heading"
+    );
+    expect(functionNameInput).toBeInTheDocument();
+    expect(functionNameInput.value).toBe("");
+    fireEvent.change(functionNameInput, {
+      target: { value: "IP" },
+    });
+    expect(functionNameInput.value).toBe("IP");
+
+    const definitionCommentTextBox = await screen.findByRole("textbox", {
+      name: "Comment",
+    });
+    expect(definitionCommentTextBox).toBeInTheDocument();
+    const definitionCommentInput = (await screen.findByTestId(
+      "function-comment-text"
+    )) as HTMLInputElement;
+    expect(definitionCommentInput.value).toBe("");
+    fireEvent.change(definitionCommentInput, {
+      target: { value: "comment" },
+    });
+    expect(definitionCommentInput.value).toBe("comment");
+
+    expect(
+      screen.getByTestId("terminology-section-Expression Editor-sub-heading")
+    ).toBeInTheDocument();
+    const typeInput = screen.getByTestId(
+      "type-selector-input"
+    ) as HTMLInputElement;
+    expect(typeInput).toBeInTheDocument();
+    expect(typeInput.value).toBe("");
+
+    fireEvent.change(typeInput, {
+      target: { value: "Timing" },
+    });
+    expect(typeInput.value).toBe("Timing");
+
+    const nameAutoComplete = screen.getByTestId("name-selector");
+    expect(nameAutoComplete).toBeInTheDocument();
+    const nameComboBox = within(nameAutoComplete).getByRole("combobox");
+    //name dropdown is populated with values based on type
+    await waitFor(() => expect(nameComboBox).toBeEnabled());
+
+    const nameDropDown = await screen.findByTestId("name-selector");
+    fireEvent.keyDown(nameDropDown, { key: "ArrowDown" });
+
+    const nameOptions = await screen.findAllByRole("option");
+    expect(nameOptions).toHaveLength(70);
+    const insertBtn = screen.getByTestId("expression-insert-btn");
+
+    expect(insertBtn).toBeInTheDocument();
+    expect(insertBtn).toBeDisabled();
+
+    fireEvent.click(nameOptions[0]);
+    expect(insertBtn).toBeEnabled();
+
+    fireEvent.click(insertBtn);
+    const definitionName = (await screen.findByTestId(
+      "function-name-text-input"
+    )) as HTMLInputElement;
+    expect(definitionName.value).toBe("IP");
+    // args
+    const argumentSectionButton = await within(argumentsSection).findByRole(
+      "button"
+    );
+    fireEvent.click(argumentSectionButton);
+
+    const argumentNameInput = (await screen.findByTestId(
+      "argument-name-input"
+    )) as HTMLInputElement;
+    expect(argumentNameInput).toBeInTheDocument();
+    expect(argumentNameInput.value).toBe("");
+    fireEvent.change(argumentNameInput, {
+      target: { value: "Test" },
+    });
+    expect(argumentNameInput.value).toBe("Test");
+
+    const dataTypeDropdown = await screen.findByTestId(
+      "arg-type-selector-input"
+    );
+    fireEvent.change(dataTypeDropdown, {
+      target: { value: "Boolean" },
+    });
+
+    const addButton = screen.getByTestId("function-argument-add-btn");
+    expect(addButton).toBeInTheDocument();
+    expect(addButton).toBeEnabled();
+    // get available dataTypes 0
+
+    fireEvent.click(addButton);
+
+    const functionArgumentTable = screen.getByTestId("function-argument-tbl");
+    expect(functionArgumentTable).toBeInTheDocument();
+    const tableRow = functionArgumentTable.querySelector("tbody").children[0];
+    expect(tableRow.children[1].textContent).toEqual("Test");
+    // submit
+    const submitButton = screen.getByTestId("function-apply-btn");
+    expect(submitButton).toBeEnabled();
+    fireEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(handleApplyFn).toHaveBeenCalledWith(
+        expect.objectContaining({
+          comment: "comment",
+          expressionValue: "after",
+          fluentFunction: true,
+          functionName: "IP",
+          functionsArguments: [
+            expect.objectContaining({
+              argumentName: "Test",
+              dataType: "Boolean",
+            }),
+          ],
+        })
+      );
+    });
+  });
 });
 
 describe("getNewExpressionsAndLines", () => {
diff --git a/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.tsx b/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.tsx
index af39a8b..37e8e0c 100644
--- a/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.tsx
+++ b/src/CqlBuilderPanel/functionsSection/functionBuilder/FunctionBuilder.tsx
@@ -222,24 +222,16 @@ export default function FunctionBuilder({
               const functionToApply = {
                 functionName: formik.values.functionName,
                 comment: formik.values.comment,
+                functionsArguments: formik.values.functionsArguments,
                 fluentFunction: formik.values.fluentFunction,
                 expressionValue: expressionEditorValue,
               };
               resetForm();
               setExpressionEditorValue("");
-              if (operation === "edit") {
-                formik.setFieldValue("functionName", "");
-                formik.setFieldValue("comment", "");
-                formik.setFieldValue("fluentFunction", true);
-                handleFunctionEdit(cqlFunction, functionToApply);
-                onClose();
-              } else {
-                handleApplyFunction(functionToApply);
-              }
+              handleApplyFunction(functionToApply);
             }}
           >
             Apply
-            {operation === "edit" ? "Save" : "Apply"}
           </Button>
         </div>
         <ConfirmationDialog