React Testing Library custom render function and async await

I have the following custom render function for my component.

It has two modes Create and Edit.

Create is synchronous and Edit is asynchronous.

The function is as follows:

    const renderComponent = async (
      scheduleId = "",
      dialogMode = DialogMode.CREATE,
      cohorts = MOCK_COHORT_LIST,
      jobs = JOB_LIST,
      availableForSchedule = domain === Domain.COHORT ? jobs : cohorts,
    ) => {
      render(
        <AddEditScheduleDialog
          cohorts={cohorts}
          jobs={jobs}
          availableForSchedule={availableForSchedule}
          scheduleToEdit={scheduleId}
          handleToggleDialog={mockToggleDialog}
          isDialogVisible={true}
          domain={domain}
        />,
        {
          wrapper: queryWrapper,
        },
      );

      if (dialogMode === DialogMode.EDIT) {
        expect(screen.getByRole("progressbar")).toBeInTheDocument();
      }

      await waitFor(() =>
        expect(screen.queryByRole("progressbar")).not.toBeInTheDocument(),
      );

      return {
        header: screen.getByRole("heading", {
          level: 1,
          name: `schedules:addEditDialog.${domain}.${dialogMode}.title`,
        }),
        name: {
          field: screen.getByTestId(`${domain}-name-field`),
          button: within(screen.getByTestId(`${domain}-name-field`)).getByRole(
            "button",
          ),
        },
        frequency: {
          field: screen.getByTestId("schedule-frequency-field"),
          input: within(
            screen.getByTestId("schedule-frequency-field"),
          ).getByRole("textbox"),
          helperText: `schedules:addEditDialog.form.frequency.helperText`,
        },
      };
    };

Sometimes I get intermittent problems finding elements on the screen. Is this because the function returns before the progressbar has been awaited?

Is there a way i can wait for everything to be rendered prior to returning the screen elements that I need?