Fix bug in form save and resume flow (#721)

* Fix bug in form save and resume flow

* Add comments

app/src/client/components/buildPage/useFormHandler.ts

@@ -24,10 +24,16 @@ export function useFormHandler(setActiveProperty: (activeProperty: string) => vo
     return stack.map((parser) => parser.getFormattedPropertyName());
   }, [stack]);
 
+  /**
+   * Pushes a new parser to the stack and updates the existing one based on user actions.
+   * @param customOptions Options for creating or updating a parser.
+   */
   const pushNewParser = useCallback((customOptions: PushParser) => {
+    // This function gets called whenever user tries to create a new form or a nested form
     const newParser = new PropertySchemaParser(customOptions.propertiesSchema, customOptions.activeProperty);
     newParser.setActiveModel(customOptions.activeModel);
 
+    // Configure the new parser
     if (customOptions.flow) {
       newParser.setFlow(customOptions.flow);
     }
@@ -40,19 +46,32 @@ export function useFormHandler(setActiveProperty: (activeProperty: string) => vo
       newParser.setUserProperties(customOptions.userProperties);
     }
 
+    // Update the parser stack to control dynamic form field rendering
     setStack((prevStack) => {
-      // When the model dropdown is updated within the same property, we need to replace the last parser in the stack
       if (customOptions.replaceLast && prevStack.length > 0) {
+        // Replace the last parser in the stack with the new parser when updating the 'select model' dropdown instead of appending to it
         return [...prevStack.slice(0, -1), newParser];
       } else {
-        // this is for rest of the cases
+        // This will be called when user creates a new nested form
         if (customOptions.formState) {
           const lastParser = prevStack[prevStack.length - 1];
-          lastParser &&
+          if (lastParser) {
+            /**
+             * Update the last parser before adding a new one for nested forms.
+             * @remarks
+             * This process involves:
+             * 1. Saving the current state of the outer form in json_str.
+             *    This retains the values when navigating back to the outer form.
+             * 2. Storing the fieldKey to identify the outer form's key.
+             *    This allows pre-population of the outer form with nested form values.
+             */
+            const lastParserActiveModelObj = lastParser.getActiveModelObj() || {};
             lastParser.setActiveModelObj({
+              ...lastParserActiveModelObj,
               json_str: customOptions.formState.values,
               fieldKey: customOptions.formState.fieldKey,
             });
+          }
         }
 
         return [...prevStack, newParser];

app/src/client/tests/UserProperty-e2e.test.tsx

@@ -1,10 +1,11 @@
 import React from 'react';
-import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { screen, waitFor, RenderResult } from '@testing-library/react';
+import { describe, it, expect, vi, beforeEach, Mock } from 'vitest';
+import { screen, waitFor } from '@testing-library/react';
 import userEvent, { UserEvent } from '@testing-library/user-event';
 import { renderInContext } from 'wasp/client/test';
 import { useQuery } from 'wasp/client/operations';
-import { mockProps } from './mocks';
+import { llmUserProperties, mockProps } from './mocks';
+import _ from 'lodash';
 
 // Types
 interface MockData {
@@ -24,7 +25,7 @@ const mockRefetch = vi.fn();
 const mockGetModels = vi.fn();
 const mockValidateForm = vi.fn();
 const mockAddUserModels = vi.fn();
-
+const mockUpdateUserModels = vi.fn();
 vi.mock('wasp/client/operations', () => ({
   useQuery: vi.fn(() => ({
     data: mockData,
@@ -34,6 +35,7 @@ vi.mock('wasp/client/operations', () => ({
   getModels: (...args: any[]) => mockGetModels(...args),
   validateForm: (...args: any[]) => mockValidateForm(...args),
   addUserModels: (...args: any[]) => mockAddUserModels(...args),
+  updateUserModels: (...args: any[]) => mockUpdateUserModels(...args),
 }));
 
 // Helper functions
@@ -343,4 +345,74 @@ describe('UserProperty Component Tests', () => {
       expect(screen.getByLabelText('Name')).toHaveValue('My Anthropic Agent');
     });
   });
+
+  it('Should call the update function correctly when the top level model is updated', async () => {
+    // The below test checks the following:
+    // 1. Click one of the existing LLM
+    // 2. Changes the api_key for that LLM, and selects the 'Add new secret' option from the dropdown
+    // 3. Creates a new secret key and ensures the addModel function is called
+    // 4. Now saves the LLM form and ensures the updateModel function is called
+
+    const user = userEvent.setup();
+    const { UserProperty } = await import('../components/buildPage/UserProperty');
+
+    // Mock useQuery to return llmUserProperties
+    const data = _.cloneDeep(llmUserProperties);
+    (useQuery as Mock).mockReturnValue({
+      data: data,
+      refetch: vi.fn().mockResolvedValue({ data: data }),
+      isLoading: false,
+    });
+
+    renderInContext(<UserProperty {...mockProps} activeProperty='llm' />);
+
+    // click the data-testid which starts with the regex "model-item-"
+    await user.click(screen.getByTestId(/model-item-.*/));
+    expect(screen.getByText('Update LLM')).toBeInTheDocument();
+
+    // Create Secret
+    await user.click(screen.getAllByRole('combobox')[0]);
+    await user.click(screen.getByText('Add new "Secret"'));
+    await waitFor(() => expect(screen.getByText('Select Secret')).toBeInTheDocument());
+    expect(screen.getByText('AzureOAIAPIKey')).toBeInTheDocument();
+    await user.type(screen.getByLabelText('Name'), 'My AzureOAIAPIKey Secret');
+    await user.type(screen.getByLabelText('API Key'), 'My Api Key');
+    mockValidateForm.mockResolvedValue({
+      name: 'My AzureOAIAPIKey Secret',
+      api_key: 'test-api-key-12345678910', // pragma: allowlist secret
+      uuid: 'test-uuid-secret-123',
+    });
+    data.push({
+      uuid: 'test-uuid-secret-123',
+      user_uuid: 'test-user-uuid-563ec0df-3ecd-4d2a',
+      type_name: 'secret',
+      model_name: 'AzureOAIAPIKey',
+      json_str: {
+        name: 'My AzureOAIAPIKey Secret',
+        api_key: 'test-api-key-12345678910', // pragma: allowlist secret
+      },
+      created_at: '2024-08-19T11:28:32.875000Z',
+      updated_at: '2024-08-19T11:28:32.875000Z',
+    });
+    // Mock useQuery to return llmUserProperties
+    (useQuery as Mock).mockReturnValue({
+      data,
+      refetch: vi.fn().mockResolvedValue({ data: data }),
+      isLoading: false,
+    });
+    // Save the secret form
+    await user.click(screen.getByRole('button', { name: 'Save' }));
+
+    // Check if the addModel function is called correctly
+    expect(mockAddUserModels).toHaveBeenCalledOnce();
+    expect(mockUpdateUserModels).not.toHaveBeenCalled();
+    expect(screen.getByText('My AzureOAIAPIKey Secret')).toBeInTheDocument();
+
+    // Save the LLM form
+    await user.click(screen.getByRole('button', { name: 'Save' }));
+
+    // Check if the updateModel function is called correctly
+    expect(mockUpdateUserModels).toHaveBeenCalledOnce();
+    expect(screen.getByText('LLM')).toBeInTheDocument();
+  });
 });