Add unit test for useFetchAgentFrameworkTraces Hook (#77)

* Add unit test for useFetchAgentFrameworkTraces Signed-off-by: Hailong Cui <[email protected]> * Add test case for unmount Signed-off-by: Hailong Cui <[email protected]> * Address review comments Signed-off-by: Hailong Cui <[email protected]> --------- Signed-off-by: Hailong Cui <[email protected]>
opensearch-project · Dec 18, 2023 · 66292ce · 66292ce
1 parent 866a679
commit 66292ce
Show file tree

Hide file tree

Showing 2 changed files with 128 additions and 3 deletions.
diff --git a/public/hooks/use_fetch_agentframework_traces.test.ts b/public/hooks/use_fetch_agentframework_traces.test.ts
@@ -0,0 +1,120 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { renderHook, act } from '@testing-library/react-hooks';
+import { useFetchAgentFrameworkTraces } from './use_fetch_agentframework_traces';
+import { createOpenSearchDashboardsReactContext } from '../../../../src/plugins/opensearch_dashboards_react/public';
+import { coreMock } from '../../../../src/core/public/mocks';
+import { HttpHandler } from '../../../../src/core/public';
+import { AbortError } from '../../../../src/plugins/data/common';
+
+describe('useFetchAgentFrameworkTraces hook', () => {
+  const traceId = 'foo';
+  const services = coreMock.createStart();
+  const { Provider } = createOpenSearchDashboardsReactContext(services);
+  const wrapper = { wrapper: Provider };
+
+  it('return undefined when trace id is not specfied', () => {
+    const { result } = renderHook(() => useFetchAgentFrameworkTraces(''));
+    expect(result.current).toMatchObject({
+      data: undefined,
+      loading: false,
+    });
+  });
+
+  it('return trace data successfully', async () => {
+    const traces = [
+      {
+        interactionId: 'test_interactionId',
+        parentInteractionId: 'test_parent_interactionId',
+        input: 'input',
+        output: 'output',
+        createTime: '',
+        origin: '',
+        traceNumber: 1,
+      },
+    ];
+
+    services.http.get.mockResolvedValueOnce(traces);
+    const { result, waitForNextUpdate } = renderHook(
+      () => useFetchAgentFrameworkTraces(traceId),
+      wrapper
+    );
+
+    await waitForNextUpdate();
+    expect(services.http.get).toHaveBeenCalledWith(
+      `/api/assistant/trace/${traceId}`,
+      expect.objectContaining({
+        signal: expect.any(Object),
+      })
+    );
+
+    expect(result.current).toEqual({
+      data: traces,
+      loading: false,
+    });
+  });
+
+  it('return error when fetch trace error happend', async () => {
+    services.http.get.mockRejectedValue(new Error('trace not found'));
+    const { result, waitForNextUpdate } = renderHook(
+      () => useFetchAgentFrameworkTraces(traceId),
+      wrapper
+    );
+
+    await waitForNextUpdate();
+
+    expect(services.http.get).toHaveBeenCalledWith(
+      `/api/assistant/trace/${traceId}`,
+      expect.objectContaining({
+        signal: expect.any(Object),
+      })
+    );
+
+    expect(result.current).toEqual({
+      data: undefined,
+      loading: false,
+      error: new Error('trace not found'),
+    });
+  });
+
+  it('abort the request when unmount', async () => {
+    const abortError = new AbortError();
+
+    const abortFn = jest.fn();
+
+    // @ts-ignore
+    global.AbortController = jest.fn(() => ({
+      abort: abortFn,
+      signal: new Object(),
+    }));
+
+    services.http.get.mockImplementation(((_path, options) => {
+      return new Promise((_resolve, reject) => {
+        if (options?.signal) {
+          options.signal.onabort = () => {
+            reject(abortError);
+          };
+        }
+      });
+    }) as HttpHandler);
+
+    const { unmount } = renderHook(() => useFetchAgentFrameworkTraces(traceId), wrapper);
+
+    act(() => {
+      unmount();
+    });
+
+    expect(services.http.get).toHaveBeenCalledWith(
+      `/api/assistant/trace/${traceId}`,
+      expect.objectContaining({
+        signal: expect.any(Object),
+      })
+    );
+
+    // expect the mock to be called
+    expect(abortFn).toBeCalledTimes(1);
+  });
+});
diff --git a/public/hooks/use_fetch_agentframework_traces.ts b/public/hooks/use_fetch_agentframework_traces.ts
@@ -23,17 +23,22 @@ export const useFetchAgentFrameworkTraces = (traceId: string) => {
     }
 
     core.services.http
-      .get<AgentFrameworkTrace[]>(`${ASSISTANT_API.TRACE}/${traceId}`)
+      .get<AgentFrameworkTrace[]>(`${ASSISTANT_API.TRACE}/${traceId}`, {
+        signal: abortController.signal,
+      })
       .then((payload) =>
         dispatch({
           type: 'success',
           payload,
         })
       )
-      .catch((error) => dispatch({ type: 'failure', error }));
+      .catch((error) => {
+        if (error.name === 'AbortError') return;
+        dispatch({ type: 'failure', error });
+      });
 
     return () => abortController.abort();
-  }, [traceId]);
+  }, [core.services.http, traceId]);
 
   return { ...state };
 };