test(ml): add unit tests for MlDataSourcePicker and DataSourceContextProvider components

- Introduced comprehensive tests for the MlDataSourcePicker component, covering rendering, data view selection, and session management functionalities.
- Added tests for the DataSourceContextProvider to ensure proper rendering based on URL parameters and error handling.
- Enhanced overall test coverage for AIOps components, improving reliability and maintainability of the codebase.

Author: KodeRad

x-pack/platform/packages/private/ml/aiops_components/src/ml_data_source_picker/ml_data_source_picker.test.tsx

@@ -0,0 +1,217 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+import { render, screen, fireEvent, act } from '@testing-library/react';
+import { __IntlProvider as IntlProvider } from '@kbn/i18n-react';
+import { MlDataSourcePicker } from './ml_data_source_picker';
+import type { MlDataSourcePickerServices } from './ml_data_source_picker';
+
+jest.mock('react-router-dom', () => ({
+  useLocation: jest.fn(() => ({ pathname: '/jobs/new_job/step/data_view', search: '' })),
+}));
+
+let capturedDataViewPickerProps: Record<string, any> = {};
+const MockDataViewPicker = (props: any) => {
+  capturedDataViewPickerProps = props;
+  return (
+    <div data-test-subj="mockDataViewPicker">
+      <span>{props.trigger?.label}</span>
+    </div>
+  );
+};
+
+jest.mock('./ml_open_session_flyout', () => ({
+  MlOpenSessionFlyout: (props: any) => {
+    return (
+      <div data-test-subj="mockOpenSessionFlyout">
+        <button onClick={props.onClose} data-test-subj="closeSessionPanel">
+          Close
+        </button>
+        <button
+          onClick={() => props.onOpenSavedSearch('saved-search-id-1')}
+          data-test-subj="openSavedSearch"
+        >
+          Open Saved Search
+        </button>
+      </div>
+    );
+  },
+}));
+
+const mockNavigateToPath = jest.fn();
+const mockGetIdsWithTitle = jest.fn().mockResolvedValue([]);
+const mockOpenEditor = jest.fn().mockResolvedValue(() => {});
+
+const buildServices = (overrides?: Partial<MlDataSourcePickerServices>): MlDataSourcePickerServices =>
+  ({
+    dataViews: { getIdsWithTitle: mockGetIdsWithTitle },
+    dataViewEditor: {
+      userPermissions: { editDataView: jest.fn(() => true) },
+    },
+    dataViewFieldEditor: { openEditor: mockOpenEditor },
+    http: { basePath: { prepend: jest.fn((p: string) => p) } },
+    application: { capabilities: {} },
+    contentManagement: { client: {} },
+    uiSettings: {},
+    ...overrides,
+  } as unknown as MlDataSourcePickerServices);
+
+const MockSavedObjectFinder = () => <div data-test-subj="mockSavedObjectFinder" />;
+
+const renderComponent = (props: { currentDataView: any; services?: MlDataSourcePickerServices }) =>
+  render(
+    <IntlProvider locale="en">
+      <MlDataSourcePicker
+        currentDataView={props.currentDataView}
+        services={props.services ?? buildServices()}
+        navigateToPath={mockNavigateToPath}
+        DataViewPickerComponent={MockDataViewPicker}
+        SavedObjectFinderComponent={MockSavedObjectFinder}
+      />
+    </IntlProvider>
+  );
+
+describe('MlDataSourcePicker', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    capturedDataViewPickerProps = {};
+  });
+
+  it('renders DataViewPicker with "Select data view" label when currentDataView is null', async () => {
+    await act(async () => {
+      renderComponent({ currentDataView: null });
+    });
+
+    expect(screen.getByTestId('mockDataViewPicker')).toBeDefined();
+    expect(screen.getByText('Select data view')).toBeDefined();
+    expect(capturedDataViewPickerProps.trigger?.label).toBe('Select data view');
+  });
+
+  it('renders DataViewPicker with the data view name when currentDataView is provided', async () => {
+    const mockDataView = {
+      id: 'dv-1',
+      getName: jest.fn(() => 'My Data View'),
+    };
+
+    await act(async () => {
+      renderComponent({ currentDataView: mockDataView });
+    });
+
+    expect(screen.getByText('My Data View')).toBeDefined();
+    expect(capturedDataViewPickerProps.trigger?.label).toBe('My Data View');
+    expect(capturedDataViewPickerProps.currentDataViewId).toBe('dv-1');
+  });
+
+  it('calls navigateToPath with ?index=<id> when onChangeDataView is triggered', async () => {
+    await act(async () => {
+      renderComponent({ currentDataView: null });
+    });
+
+    await act(async () => {
+      capturedDataViewPickerProps.onChangeDataView('test-index-id');
+    });
+
+    expect(mockNavigateToPath).toHaveBeenCalledWith(
+      '/jobs/new_job/step/data_view?index=test-index-id'
+    );
+  });
+
+  it('renders MlOpenSessionFlyout when "Open Discover session" button is clicked', async () => {
+    await act(async () => {
+      renderComponent({ currentDataView: null });
+    });
+
+    expect(screen.queryByTestId('mockOpenSessionFlyout')).toBeNull();
+
+    await act(async () => {
+      fireEvent.click(screen.getByTestId('mlOpenDiscoverSessionButton'));
+    });
+
+    expect(screen.getByTestId('mockOpenSessionFlyout')).toBeDefined();
+  });
+
+  it('onOpenSavedSearch calls navigateToPath with ?savedSearchId=<id> and hides the flyout', async () => {
+    await act(async () => {
+      renderComponent({ currentDataView: null });
+    });
+
+    await act(async () => {
+      fireEvent.click(screen.getByTestId('mlOpenDiscoverSessionButton'));
+    });
+
+    expect(screen.getByTestId('mockOpenSessionFlyout')).toBeDefined();
+
+    await act(async () => {
+      fireEvent.click(screen.getByTestId('openSavedSearch'));
+    });
+
+    expect(mockNavigateToPath).toHaveBeenCalledWith(
+      '/jobs/new_job/step/data_view?savedSearchId=saved-search-id-1'
+    );
+    expect(screen.queryByTestId('mockOpenSessionFlyout')).toBeNull();
+  });
+
+  it('onDataViewCreated navigates and refreshes data views when the new view has an id', async () => {
+    const refreshedViews = [{ id: 'new-dv', title: 'New DV' }];
+    mockGetIdsWithTitle.mockResolvedValueOnce([]).mockResolvedValueOnce(refreshedViews);
+
+    await act(async () => {
+      renderComponent({ currentDataView: null });
+    });
+
+    await act(async () => {
+      await capturedDataViewPickerProps.onDataViewCreated({ id: 'new-dv-id' });
+    });
+
+    expect(mockNavigateToPath).toHaveBeenCalledWith(
+      '/jobs/new_job/step/data_view?index=new-dv-id'
+    );
+    expect(mockGetIdsWithTitle).toHaveBeenCalledTimes(2);
+  });
+
+  it('onAddField is defined and calls dataViewFieldEditor.openEditor when canEditDataView=true and currentDataView is set', async () => {
+    const mockDataView = {
+      id: 'dv-1',
+      getName: jest.fn(() => 'My Data View'),
+    };
+
+    await act(async () => {
+      renderComponent({ currentDataView: mockDataView });
+    });
+
+    expect(capturedDataViewPickerProps.onAddField).toBeDefined();
+
+    await act(async () => {
+      await capturedDataViewPickerProps.onAddField();
+    });
+
+    expect(mockOpenEditor).toHaveBeenCalledWith({
+      ctx: { dataView: mockDataView },
+      onSave: expect.any(Function),
+    });
+  });
+
+  it('onAddField is undefined when canEditDataView=false', async () => {
+    const mockDataView = {
+      id: 'dv-1',
+      getName: jest.fn(() => 'My Data View'),
+    };
+
+    const services = buildServices({
+      dataViewEditor: {
+        userPermissions: { editDataView: jest.fn(() => false) },
+      },
+    } as any);
+
+    await act(async () => {
+      renderComponent({ currentDataView: mockDataView, services });
+    });
+
+    expect(capturedDataViewPickerProps.onAddField).toBeUndefined();
+  });
+});

x-pack/platform/plugins/shared/ml/public/application/contexts/ml/data_source_context.test.tsx

@@ -0,0 +1,160 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
+import { __IntlProvider as IntlProvider } from '@kbn/i18n-react';
+import { DataSourceContextProvider } from './data_source_context';
+import { useMlKibana } from '../kibana';
+
+jest.mock('../kibana', () => ({
+  useMlKibana: jest.fn(),
+}));
+
+jest.mock('../../util/index_utils', () => ({
+  getDataViewAndSavedSearchCallback: jest.fn(() => async (id: string) => ({
+    dataView: { id, title: 'mock-saved-search-data-view' },
+    savedSearch: { id: 'mock-saved-search' },
+  })),
+}));
+
+jest.mock('../../jobs/new_job/utils/new_job_utils', () => ({
+  createSearchItems: jest.fn(() => ({ combinedQuery: { match_all: {} } })),
+}));
+
+const mockLocationSearch = jest.fn(() => '');
+jest.mock('react-router-dom', () => ({
+  useLocation: () => ({ pathname: '/', search: mockLocationSearch() }),
+}));
+
+const mockGet = jest.fn();
+const mockGetDefaultDataView = jest.fn();
+const mockGetDataViewAndSavedSearch = jest.fn();
+
+const buildKibanaMock = () => ({
+  services: {
+    data: {
+      dataViews: {
+        get: mockGet,
+        getDefaultDataView: mockGetDefaultDataView,
+      },
+    },
+    savedSearch: mockGetDataViewAndSavedSearch,
+    uiSettings: {},
+  },
+});
+
+const renderProvider = (children = <div data-test-subj="child-content">Hello</div>) =>
+  render(
+    <IntlProvider locale="en">
+      <DataSourceContextProvider>{children}</DataSourceContextProvider>
+    </IntlProvider>
+  );
+
+describe('DataSourceContextProvider', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    (useMlKibana as jest.Mock).mockReturnValue(buildKibanaMock());
+  });
+
+  it('renders children when default data view is found (no URL params)', async () => {
+    mockLocationSearch.mockReturnValue('');
+    mockGetDefaultDataView.mockResolvedValue({
+      id: 'default-dv',
+      title: 'Default Data View',
+    });
+
+    renderProvider();
+
+    await waitFor(() => {
+      expect(screen.getByTestId('child-content')).toBeInTheDocument();
+    });
+
+    expect(mockGetDefaultDataView).toHaveBeenCalled();
+  });
+
+  it('renders children when index URL param is present', async () => {
+    mockLocationSearch.mockReturnValue('?index=my-index-id');
+    mockGet.mockResolvedValue({
+      id: 'my-index-id',
+      title: 'My Index',
+    });
+
+    renderProvider();
+
+    await waitFor(() => {
+      expect(screen.getByTestId('child-content')).toBeInTheDocument();
+    });
+
+    expect(mockGet).toHaveBeenCalledWith('my-index-id');
+  });
+
+  it('renders children when savedSearchId URL param is present', async () => {
+    const { getDataViewAndSavedSearchCallback } = jest.requireMock('../../util/index_utils');
+    getDataViewAndSavedSearchCallback.mockReturnValue(async (id: string) => ({
+      dataView: { id: 'dv-from-saved-search', title: 'From Saved Search' },
+      savedSearch: { id },
+    }));
+
+    mockLocationSearch.mockReturnValue('?savedSearchId=my-saved-search');
+
+    renderProvider();
+
+    await waitFor(() => {
+      expect(screen.getByTestId('child-content')).toBeInTheDocument();
+    });
+  });
+
+  it('shows error state when resolveDataSource throws', async () => {
+    mockLocationSearch.mockReturnValue('?index=bad-index');
+    mockGet.mockRejectedValue(new Error('Data view not found'));
+
+    renderProvider();
+
+    await waitFor(() => {
+      expect(
+        screen.getByText('Unable to fetch data view or saved Discover session')
+      ).toBeInTheDocument();
+    });
+
+    expect(screen.queryByTestId('child-content')).not.toBeInTheDocument();
+  });
+
+  it('renders null initially before the async resolve completes', async () => {
+    mockLocationSearch.mockReturnValue('');
+    let resolvePromise: (value: any) => void;
+    mockGetDefaultDataView.mockReturnValue(
+      new Promise((resolve) => {
+        resolvePromise = resolve;
+      })
+    );
+
+    const { container } = renderProvider();
+
+    expect(container.firstChild).toBeNull();
+
+    await waitFor(() => {
+      resolvePromise!({ id: 'default-dv', title: 'Default' });
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId('child-content')).toBeInTheDocument();
+    });
+  });
+
+  it('shows error when dataViewId is an empty string', async () => {
+    mockLocationSearch.mockReturnValue('?index=');
+
+    renderProvider();
+
+    await waitFor(() => {
+      expect(
+        screen.getByText('Unable to fetch data view or saved Discover session')
+      ).toBeInTheDocument();
+    });
+  });
+});