From d97ce5b346d386e057dfbd131c8f7539661ca1f7 Mon Sep 17 00:00:00 2001
From: Anan Zhuang <ananzh@amazon.com>
Date: Wed, 16 Oct 2024 18:06:07 +0000
Subject: [PATCH] fix the test by adding the getUpdates$ method to the mock
 queryString object

Signed-off-by: Anan Zhuang <ananzh@amazon.com>
---
 .../public/ui/dataset_selector/index.test.tsx | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/plugins/data/public/ui/dataset_selector/index.test.tsx b/src/plugins/data/public/ui/dataset_selector/index.test.tsx
index 0dd456711c37..db78461b96c1 100644
--- a/src/plugins/data/public/ui/dataset_selector/index.test.tsx
+++ b/src/plugins/data/public/ui/dataset_selector/index.test.tsx
@@ -5,6 +5,7 @@
 
 import React from 'react';
 import { mount } from 'enzyme';
+import { act } from 'react-dom/test-utils';
 import { DatasetSelector as ConnectedDatasetSelector } from './index';
 import { DatasetSelector } from './dataset_selector';
 import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
@@ -19,6 +20,8 @@ jest.mock('./dataset_selector', () => ({
 }));
 
 describe('ConnectedDatasetSelector', () => {
+  const mockSubscribe = jest.fn();
+  const mockUnsubscribe = jest.fn();
   const mockQueryString = {
     getQuery: jest.fn().mockReturnValue({}),
     getDefaultQuery: jest.fn().mockReturnValue({}),
@@ -27,12 +30,14 @@ describe('ConnectedDatasetSelector', () => {
     getDatasetService: jest.fn().mockReturnValue({
       addRecentDataset: jest.fn(),
     }),
+    getUpdates$: jest.fn().mockReturnValue({
+      subscribe: mockSubscribe.mockReturnValue({ unsubscribe: mockUnsubscribe }),
+    }),
   };
   const mockOnSubmit = jest.fn();
   const mockServices = {
     data: {
       query: {
-        // @ts-ignore
         queryString: mockQueryString,
       },
     },
@@ -40,6 +45,7 @@ describe('ConnectedDatasetSelector', () => {
 
   beforeEach(() => {
     (useOpenSearchDashboards as jest.Mock).mockReturnValue({ services: mockServices });
+    jest.clearAllMocks();
   });
 
   it('should render DatasetSelector with correct props', () => {
@@ -66,10 +72,23 @@ describe('ConnectedDatasetSelector', () => {
     ) => void;
 
     const newDataset: Dataset = { id: 'test', title: 'Test Dataset', type: 'test' };
-    setSelectedDataset(newDataset);
+    act(() => {
+      setSelectedDataset(newDataset);
+    });
 
     expect(mockQueryString.getInitialQueryByDataset).toHaveBeenCalledTimes(1);
     expect(mockQueryString.setQuery).toHaveBeenCalledTimes(1);
     expect(mockOnSubmit).toHaveBeenCalledTimes(1);
   });
+
+  it('should subscribe to queryString.getUpdates$ and unsubscribe on unmount', () => {
+    const wrapper = mount(<ConnectedDatasetSelector onSubmit={mockOnSubmit} />);
+
+    expect(mockQueryString.getUpdates$).toHaveBeenCalledTimes(1);
+    expect(mockSubscribe).toHaveBeenCalledTimes(1);
+
+    wrapper.unmount();
+
+    expect(mockUnsubscribe).toHaveBeenCalledTimes(1);
+  });
 });