From 34f8d2d62ba29f7aa85ff18afbf09c4f7e1e16c0 Mon Sep 17 00:00:00 2001
From: Danny Banks <djb@amazon.com>
Date: Fri, 8 Nov 2024 16:24:56 -0800
Subject: [PATCH] chore(ai): fixing ai test coverage (#6037)

---
 .../AIConversation/views/default/Form.tsx     |  2 +-
 .../views/default/PromptList.tsx              |  2 +-
 .../default/__tests__/PromptList.test.tsx     | 71 +++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 packages/react-ai/src/components/AIConversation/views/default/__tests__/PromptList.test.tsx

diff --git a/packages/react-ai/src/components/AIConversation/views/default/Form.tsx b/packages/react-ai/src/components/AIConversation/views/default/Form.tsx
index 3cc7c66653f..3b1f3e155dc 100644
--- a/packages/react-ai/src/components/AIConversation/views/default/Form.tsx
+++ b/packages/react-ai/src/components/AIConversation/views/default/Form.tsx
@@ -49,7 +49,7 @@ const FormWrapper = ({
   }
 };
 
-export const Form: NonNullable<ControlsContextProps['Form']> = ({
+export const Form: Required<ControlsContextProps>['Form'] = ({
   setInput,
   input,
   handleSubmit,
diff --git a/packages/react-ai/src/components/AIConversation/views/default/PromptList.tsx b/packages/react-ai/src/components/AIConversation/views/default/PromptList.tsx
index fab975e4034..c79b07d85d0 100644
--- a/packages/react-ai/src/components/AIConversation/views/default/PromptList.tsx
+++ b/packages/react-ai/src/components/AIConversation/views/default/PromptList.tsx
@@ -4,7 +4,7 @@ import { Flex, Button } from '@aws-amplify/ui-react';
 import { ControlsContextProps } from '../../context/ControlsContext';
 import { ComponentClassName } from '@aws-amplify/ui';
 
-export const PromptList: ControlsContextProps['PromptList'] = ({
+export const PromptList: Required<ControlsContextProps>['PromptList'] = ({
   setInput,
   suggestedPrompts = [],
 }) => {
diff --git a/packages/react-ai/src/components/AIConversation/views/default/__tests__/PromptList.test.tsx b/packages/react-ai/src/components/AIConversation/views/default/__tests__/PromptList.test.tsx
new file mode 100644
index 00000000000..70c2e80ff06
--- /dev/null
+++ b/packages/react-ai/src/components/AIConversation/views/default/__tests__/PromptList.test.tsx
@@ -0,0 +1,71 @@
+import React from 'react';
+import { render, fireEvent } from '@testing-library/react';
+import { PromptList } from '../PromptList';
+import { ComponentClassName } from '@aws-amplify/ui';
+import { ConversationInputContext } from '../../../context';
+
+describe('PromptList', () => {
+  const mockSetInput = jest.fn<
+    ReturnType<Required<ConversationInputContext>['setInput']>,
+    Parameters<Required<ConversationInputContext>['setInput']>
+  >();
+
+  it('renders without crashing', () => {
+    const { container } = render(
+      <PromptList setInput={mockSetInput} suggestedPrompts={[]} />
+    );
+    expect(container).toBeTruthy();
+  });
+
+  it('renders suggested prompts', () => {
+    const suggestedPrompts = [
+      { inputText: 'Prompt 1', component: 'Prompt 1' },
+      { inputText: 'Prompt 2', component: 'Prompt 2' },
+    ];
+    const { getByText } = render(
+      <PromptList setInput={mockSetInput} suggestedPrompts={suggestedPrompts} />
+    );
+
+    expect(getByText('Prompt 1')).toBeInTheDocument();
+    expect(getByText('Prompt 2')).toBeInTheDocument();
+  });
+
+  it('calls setInput with the correct prompt when a button is clicked', () => {
+    const suggestedPrompts = [
+      { inputText: 'Prompt 1', component: 'Prompt 1' },
+      { inputText: 'Prompt 2', component: 'Prompt 2' },
+    ];
+    const { getByText } = render(
+      <PromptList setInput={mockSetInput} suggestedPrompts={suggestedPrompts} />
+    );
+
+    fireEvent.click(getByText('Prompt 1'));
+    expect(mockSetInput).toHaveBeenCalledWith(expect.any(Function));
+
+    // In the PromptList component we are using an updater function
+    // to update the input state with setInput. In this test we are
+    // mocking setInput, so here we are grabbing what the mock was called with
+    // (a function) and calling that function to ensure
+    // this component would update the input state properly
+    const setInputCall = mockSetInput.mock.calls[0][0] as Function;
+    const result = setInputCall({ text: 'previous text' });
+    expect(result).toEqual({ text: 'Prompt 1' });
+  });
+
+  it('applies the correct CSS class to the buttons', () => {
+    const suggestedPrompts = [{ inputText: 'Prompt 1', component: 'Prompt 1' }];
+    const { getByText } = render(
+      <PromptList setInput={mockSetInput} suggestedPrompts={suggestedPrompts} />
+    );
+
+    const button = getByText('Prompt 1');
+    expect(button).toHaveClass(ComponentClassName.AIConversationPrompt);
+  });
+
+  it('renders nothing when suggestedPrompts is empty', () => {
+    const { container } = render(
+      <PromptList setInput={mockSetInput} suggestedPrompts={[]} />
+    );
+    expect(container.firstChild).toBeEmptyDOMElement();
+  });
+});