From f86eb0deeb88d631e88492fed746577df9886a1a Mon Sep 17 00:00:00 2001
From: Trung Le <58827450+ltrung@users.noreply.github.com>
Date: Fri, 5 Apr 2024 10:16:05 -0700
Subject: [PATCH] Add await for toggleContentShare (#946)

---
 src/providers/ContentShareProvider/index.tsx  |  8 +--
 .../ContentShareProvider/index.test.tsx       | 55 +++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 tst/providers/ContentShareProvider/index.test.tsx

diff --git a/src/providers/ContentShareProvider/index.tsx b/src/providers/ContentShareProvider/index.tsx
index 2fb162d1..26abacd2 100644
--- a/src/providers/ContentShareProvider/index.tsx
+++ b/src/providers/ContentShareProvider/index.tsx
@@ -128,11 +128,11 @@ export const ContentShareProvider: React.FC<
         audioVideo.stopContentShare();
       } else {
         if (source && typeof source === 'string') {
-          audioVideo.startContentShareFromScreenCapture(source);
-        } else if (source instanceof MediaStream) {
-          audioVideo.startContentShare(source);
+          await audioVideo.startContentShareFromScreenCapture(source);
+        } else if (source && source instanceof MediaStream) {
+          await audioVideo.startContentShare(source);
         } else {
-          audioVideo.startContentShareFromScreenCapture();
+          await audioVideo.startContentShareFromScreenCapture();
         }
         dispatch({ type: ContentActionType.STARTING });
       }
diff --git a/tst/providers/ContentShareProvider/index.test.tsx b/tst/providers/ContentShareProvider/index.test.tsx
new file mode 100644
index 00000000..2a1bc752
--- /dev/null
+++ b/tst/providers/ContentShareProvider/index.test.tsx
@@ -0,0 +1,55 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import React from 'react';
+import { renderHook, act } from '@testing-library/react';
+import {
+  ContentShareProvider,
+  useContentShareControls,
+} from '../../../src/providers/ContentShareProvider';
+
+// Mock audioVideo object
+const mockAudioVideo = {
+    startContentShareFromScreenCapture: jest.fn(),
+    startContentShare: jest.fn(),
+    stopContentShare: jest.fn(),
+    addObserver: jest.fn(),
+    removeObserver: jest.fn(),
+    addContentShareObserver: jest.fn(),
+    removeContentShareObserver: jest.fn()
+ };
+ // Mock the module containing useAudioVideo hook
+jest.mock('../../../src/providers/AudioVideoProvider', () => ({
+    useAudioVideo: () => mockAudioVideo
+}));
+
+describe('toggleContentShare', () => {
+  it('Should call startContentShareFromScreenCapture', async () => {
+    const { result } = renderHook(() => useContentShareControls(), {
+      wrapper: ({ children }) => (
+        <ContentShareProvider>{children}</ContentShareProvider>
+      ),
+    });
+    await act(async () => {
+      await result.current.toggleContentShare();
+    });
+    expect(mockAudioVideo.startContentShareFromScreenCapture).toHaveBeenCalledTimes(1);
+  });
+
+  it('Can catch error from startContentShareFromScreenCapture', async () => {
+    mockAudioVideo.startContentShareFromScreenCapture = jest.fn(() => Promise.reject(new Error('startContentShare error')));
+    const { result } = renderHook(() => useContentShareControls(), {
+      wrapper: ({ children }) => (
+        <ContentShareProvider>{children}</ContentShareProvider>
+      ),
+    });
+    await expect(async () => { 
+      await act(async () => {
+        await result.current.toggleContentShare();
+    });
+    }).rejects.toThrow('startContentShare error');
+    // expect(() => {
+    //     result.current.toggleContentShare();
+    //   }).toThrow('startContentShare error');
+  });
+});
\ No newline at end of file