From a1502f8f2f00d7af3c13f81ac84fb41dd1250b72 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Mon, 3 Jun 2024 09:28:57 +0000
Subject: [PATCH] Refactor with getDefaultDataSourceId$ (#197)

* Refactor with getDefaultDataSourceId$

Signed-off-by: Lin Wang <wonglam@amazon.com>

* Update change log

Signed-off-by: Lin Wang <wonglam@amazon.com>

---------

Signed-off-by: Lin Wang <wonglam@amazon.com>
(cherry picked from commit ec49168d4e896d9abcde85b583f6aeec44994542)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md
---
 public/services/__tests__/data_source_service.test.ts |  4 ++--
 public/services/data_source_service.ts                |  4 ++--
 public/types.ts                                       | 10 +---------
 3 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/public/services/__tests__/data_source_service.test.ts b/public/services/__tests__/data_source_service.test.ts
index 380dbae2..d1db20dc 100644
--- a/public/services/__tests__/data_source_service.test.ts
+++ b/public/services/__tests__/data_source_service.test.ts
@@ -8,7 +8,7 @@ import { first } from 'rxjs/operators';
 
 import { uiSettingsServiceMock } from '../../../../../src/core/public/mocks';
 import { DataSourceOption } from '../../../../../src/plugins/data_source_management/public/components/data_source_menu/types';
-import { DataSourceManagementPluginSetup } from '../../types';
+import { DataSourceManagementPluginSetup } from '../../../../../src/plugins/data_source_management/public';
 import { DataSourceService } from '../data_source_service';
 
 const setup = (options?: {
@@ -24,10 +24,10 @@ const setup = (options?: {
     dataSourceSelection: {
       getSelection$: () => dataSourceSelection$,
     },
+    getDefaultDataSourceId$: jest.fn(() => defaultDataSourceSelection$),
   };
   const dataSource = new DataSourceService();
   const defaultDataSourceSelection$ = new BehaviorSubject(options?.defaultDataSourceId ?? null);
-  uiSettings.get$.mockReturnValue(defaultDataSourceSelection$);
   const setupResult = dataSource.setup({
     uiSettings,
     dataSourceManagement:
diff --git a/public/services/data_source_service.ts b/public/services/data_source_service.ts
index 8ddb8fad..f89565e4 100644
--- a/public/services/data_source_service.ts
+++ b/public/services/data_source_service.ts
@@ -8,7 +8,7 @@ import { first, map } from 'rxjs/operators';
 
 import type { IUiSettingsClient } from '../../../../src/core/public';
 import type { DataSourceOption } from '../../../../src/plugins/data_source_management/public/components/data_source_menu/types';
-import type { DataSourceManagementPluginSetup } from '../types';
+import { DataSourceManagementPluginSetup } from '../../../../src/plugins/data_source_management/public';
 
 export enum DataSourceIdFrom {
   UiSettings,
@@ -83,7 +83,7 @@ export class DataSourceService {
   getDataSourceId$() {
     return combineLatest([
       this.dataSourceId$,
-      this.uiSettings?.get$('defaultDataSource', null) ?? of(null),
+      this.dataSourceManagement?.getDefaultDataSourceId$?.(this.uiSettings) ?? of(null),
     ]).pipe(
       map(([selectedDataSourceId, defaultDataSourceId]) => {
         if (selectedDataSourceId !== null) {
diff --git a/public/types.ts b/public/types.ts
index 02101e14..9f92ab1a 100644
--- a/public/types.ts
+++ b/public/types.ts
@@ -3,21 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { BehaviorSubject } from 'rxjs';
 import { DashboardStart } from '../../../src/plugins/dashboard/public';
 import { EmbeddableSetup, EmbeddableStart } from '../../../src/plugins/embeddable/public';
 import { IMessage, ISuggestedAction } from '../common/types/chat_saved_object_attributes';
 import { IChatContext } from './contexts/chat_context';
 import { MessageContentProps } from './tabs/chat/messages/message_content';
 import { DataSourceServiceContract, IncontextInsightRegistry } from './services';
-import { DataSourceOption } from '../../../src/plugins/data_source_management/public';
-
-// TODO: should replace from DataSourceManagementPluginSetup in DSM plugin after data selection merged
-export interface DataSourceManagementPluginSetup {
-  dataSourceSelection?: {
-    getSelection$: () => BehaviorSubject<Map<string, DataSourceOption[]>>;
-  };
-}
+import { DataSourceManagementPluginSetup } from '../../../src/plugins/data_source_management/public';
 
 export interface RenderProps {
   props: MessageContentProps;