superset2/superset-frontend/src/SqlLab/components/SaveDatasetModal/SaveDatasetModal.test.tsx

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
import * as reactRedux from 'react-redux';
import { act } from 'react';
import {
  cleanup,
  fireEvent,
  render,
  screen,
  userEvent,
  waitFor,
} from 'spec/helpers/testing-library';
import fetchMock from 'fetch-mock';
import { SaveDatasetModal } from 'src/SqlLab/components/SaveDatasetModal';
import { createDatasource } from 'src/SqlLab/actions/sqlLab';
import { user, testQuery, mockdatasets } from 'src/SqlLab/fixtures';
import { FeatureFlag, SupersetClient } from '@superset-ui/core';

const mockedProps = {
  visible: true,
  onHide: () => {},
  buttonTextOnSave: 'Save',
  buttonTextOnOverwrite: 'Overwrite',
  datasource: testQuery,
};

fetchMock.get('glob:*/api/v1/dataset/?*', {
  result: mockdatasets,
  dataset_count: 3,
});

jest.useFakeTimers({ advanceTimers: true });

// Mock the user
const useSelectorMock = jest.spyOn(reactRedux, 'useSelector');
beforeEach(() => {
  useSelectorMock.mockClear();
  cleanup();
});

// Mock the createDatasource action
const useDispatchMock = jest.spyOn(reactRedux, 'useDispatch');
jest.mock('src/SqlLab/actions/sqlLab', () => ({
  createDatasource: jest.fn(),
}));
jest.mock('src/explore/exploreUtils/formData', () => ({
  postFormData: jest.fn(),
}));
jest.mock('src/utils/cachedSupersetGet', () => ({
  ...jest.requireActual('src/utils/cachedSupersetGet'),
  clearDatasetCache: jest.fn(),
}));

// eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks
describe('SaveDatasetModal', () => {
  test('renders a "Save as new" field', () => {
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    const saveRadioBtn = screen.getByRole('radio', {
      name: /save as new/i,
    });

    const fieldLabel = screen.getByText(/save as new/i);
    const inputField = screen.getByRole('textbox');
    const inputFieldText = screen.getByDisplayValue(/unimportant/i);

    expect(saveRadioBtn).toBeInTheDocument();
    expect(fieldLabel).toBeInTheDocument();
    expect(inputField).toBeInTheDocument();
    expect(inputFieldText).toBeInTheDocument();
  });

  test('renders an "Overwrite existing" field', () => {
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    const overwriteRadioBtn = screen.getByRole('radio', {
      name: /overwrite existing/i,
    });
    const fieldLabel = screen.getByText(/overwrite existing/i);
    const inputField = screen.getByRole('combobox');
    const placeholderText = screen.getByText(/select or type dataset name/i);

    expect(overwriteRadioBtn).toBeInTheDocument();
    expect(fieldLabel).toBeInTheDocument();
    expect(inputField).toBeInTheDocument();
    expect(placeholderText).toBeInTheDocument();
  });

  test('renders a close button', () => {
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    expect(screen.getByRole('button', { name: /close/i })).toBeInTheDocument();
  });

  test('renders a save button when "Save as new" is selected', () => {
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    // "Save as new" is selected when the modal opens by default
    expect(screen.getByRole('button', { name: /save/i })).toBeInTheDocument();
  });

  test('renders an overwrite button when "Overwrite existing" is selected', () => {
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    // Click the overwrite radio button to reveal the overwrite confirmation and back buttons
    const overwriteRadioBtn = screen.getByRole('radio', {
      name: /overwrite existing/i,
    });
    userEvent.click(overwriteRadioBtn);

    expect(
      screen.getByRole('button', { name: /overwrite/i }),
    ).toBeInTheDocument();
  });

  test('renders the overwrite button as disabled until an existing dataset is selected', async () => {
    useSelectorMock.mockReturnValue({ ...user });
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    // Click the overwrite radio button
    const overwriteRadioBtn = screen.getByRole('radio', {
      name: /overwrite existing/i,
    });
    await userEvent.click(overwriteRadioBtn);

    // Overwrite confirmation button should be disabled at this point
    const overwriteConfirmationBtn = screen.getByRole('button', {
      name: /overwrite/i,
    });
    expect(overwriteConfirmationBtn).toBeDisabled();

    // Click the overwrite select component
    const select = screen.getByRole('combobox', { name: /existing dataset/i })!;
    await userEvent.click(select);

    // Advance timers to flush debounced fetches in AsyncSelect
    await act(async () => {
      jest.runAllTimers();
    });

    await waitFor(() => {
      const loading = screen.queryByText('Loading...');
      expect(loading === null || !loading.checkVisibility()).toBe(true);
    });

    // Select the first "existing dataset" from the listbox
    const option = screen.getAllByText('coolest table 0')[1];
    await userEvent.click(option);

    // Overwrite button should now be enabled
    expect(overwriteConfirmationBtn).toBeEnabled();
  });

  test('renders a confirm overwrite screen when overwrite is clicked', async () => {
    useSelectorMock.mockReturnValue({ ...user });
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    // Click the overwrite radio button
    const overwriteRadioBtn = screen.getByRole('radio', {
      name: /overwrite existing/i,
    });
    await userEvent.click(overwriteRadioBtn);

    // Click the overwrite select component
    const select = screen.getByRole('combobox', { name: /existing dataset/i });
    await userEvent.click(select);

    // Advance timers to flush debounced fetches in AsyncSelect
    await act(async () => {
      jest.runAllTimers();
    });

    await waitFor(() => {
      const loading = screen.queryByText('Loading...');
      expect(loading === null || !loading.checkVisibility()).toBe(true);
    });

    // Select the first "existing dataset" from the listbox
    const option = screen.getAllByText('coolest table 0')[1];
    await userEvent.click(option);

    // Click the overwrite button to access the confirmation screen
    const overwriteConfirmationBtn = screen.getByRole('button', {
      name: /overwrite/i,
    });
    await userEvent.click(overwriteConfirmationBtn);

    // Overwrite screen text
    expect(screen.getByText(/save or overwrite dataset/i)).toBeInTheDocument();
    expect(
      screen.getByText(/are you sure you want to overwrite this dataset\?/i),
    ).toBeInTheDocument();
    // Overwrite screen buttons
    expect(screen.getByRole('button', { name: /close/i })).toBeInTheDocument();
    expect(screen.getByRole('button', { name: /back/i })).toBeInTheDocument();
    expect(
      screen.getByRole('button', { name: /overwrite/i }),
    ).toBeInTheDocument();
  });

  test('sends the schema when creating the dataset', async () => {
    const dummyDispatch = jest.fn().mockResolvedValue({});
    useDispatchMock.mockReturnValue(dummyDispatch);
    useSelectorMock.mockReturnValue({ ...user });

    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    const inputFieldText = screen.getByDisplayValue(/unimportant/i);
    fireEvent.change(inputFieldText, { target: { value: 'my dataset' } });

    const saveConfirmationBtn = screen.getByRole('button', {
      name: /save/i,
    });
    userEvent.click(saveConfirmationBtn);

    expect(createDatasource).toHaveBeenCalledWith({
      datasourceName: 'my dataset',
      dbId: 1,
      catalog: null,
      schema: 'main',
      sql: 'SELECT *',
      templateParams: undefined,
    });
  });

  test('sends the catalog when creating the dataset', async () => {
    const dummyDispatch = jest.fn().mockResolvedValue({});
    useDispatchMock.mockReturnValue(dummyDispatch);
    useSelectorMock.mockReturnValue({ ...user });

    render(
      <SaveDatasetModal
        {...mockedProps}
        datasource={{ ...mockedProps.datasource, catalog: 'public' }}
      />,
      { useRedux: true },
    );

    const inputFieldText = screen.getByDisplayValue(/unimportant/i);
    fireEvent.change(inputFieldText, { target: { value: 'my dataset' } });

    const saveConfirmationBtn = screen.getByRole('button', {
      name: /save/i,
    });
    userEvent.click(saveConfirmationBtn);

    expect(createDatasource).toHaveBeenCalledWith({
      datasourceName: 'my dataset',
      dbId: 1,
      catalog: 'public',
      schema: 'main',
      sql: 'SELECT *',
      templateParams: undefined,
    });
  });

  test('does not renders a checkbox button when template processing is disabled', () => {
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });
    expect(screen.queryByRole('checkbox')).not.toBeInTheDocument();
  });

  test('renders a checkbox button when template processing is enabled', () => {
    // @ts-expect-error
    global.featureFlags = {
      [FeatureFlag.EnableTemplateProcessing]: true,
    };
    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });
    expect(screen.getByRole('checkbox')).toBeInTheDocument();
  });

  test('correctly includes template parameters when template processing is enabled', () => {
    // @ts-expect-error
    global.featureFlags = {
      [FeatureFlag.EnableTemplateProcessing]: true,
    };
    const propsWithTemplateParam = {
      ...mockedProps,
      datasource: {
        ...testQuery,
        templateParams: JSON.stringify({ my_param: 12 }),
      },
    };
    render(<SaveDatasetModal {...propsWithTemplateParam} />, {
      useRedux: true,
    });
    const inputFieldText = screen.getByDisplayValue(/unimportant/i);
    fireEvent.change(inputFieldText, { target: { value: 'my dataset' } });

    userEvent.click(screen.getByRole('checkbox'));

    const saveConfirmationBtn = screen.getByRole('button', {
      name: /save/i,
    });
    userEvent.click(saveConfirmationBtn);

    expect(createDatasource).toHaveBeenCalledWith({
      datasourceName: 'my dataset',
      dbId: 1,
      catalog: null,
      schema: 'main',
      sql: 'SELECT *',
      templateParams: JSON.stringify({ my_param: 12 }),
    });
  });

  test('correctly excludes template parameters when template processing is enabled', () => {
    // @ts-expect-error
    global.featureFlags = {
      [FeatureFlag.EnableTemplateProcessing]: true,
    };
    const propsWithTemplateParam = {
      ...mockedProps,
      datasource: {
        ...testQuery,
        templateParams: JSON.stringify({ my_param: 12 }),
      },
    };
    render(<SaveDatasetModal {...propsWithTemplateParam} />, {
      useRedux: true,
    });
    const inputFieldText = screen.getByDisplayValue(/unimportant/i);
    fireEvent.change(inputFieldText, { target: { value: 'my dataset' } });

    userEvent.click(screen.getByRole('checkbox'));

    const saveConfirmationBtn = screen.getByRole('button', {
      name: /save/i,
    });
    userEvent.click(saveConfirmationBtn);

    expect(createDatasource).toHaveBeenCalledWith({
      datasourceName: 'my dataset',
      dbId: 1,
      catalog: null,
      schema: 'main',
      sql: 'SELECT *',
      templateParams: undefined,
    });
  });

  const setupOverwriteFlow = async () => {
    // Select the "Overwrite existing" radio
    await userEvent.click(
      screen.getByRole('radio', { name: /overwrite existing/i }),
    );
    // Open the select to load existing-dataset options
    await userEvent.click(
      screen.getByRole('combobox', { name: /existing dataset/i }),
    );
    // Advance timers to flush debounced fetches in AsyncSelect
    await act(async () => {
      jest.runAllTimers();
    });
    // Wait for the loading indicator to clear
    await waitFor(() => {
      const loading = screen.queryByText('Loading...');
      expect(loading === null || !loading.checkVisibility()).toBe(true);
    });
    // Pick an existing dataset (use the listbox item, not the input mirror)
    const options = await screen.findAllByText('coolest table 0');
    await userEvent.click(options[1]);
    // First overwrite click → confirmation screen
    await userEvent.click(screen.getByRole('button', { name: /overwrite/i }));
    // Wait for the confirmation screen to render
    await screen.findByText(/are you sure you want to overwrite this dataset/i);
    // Second overwrite click → triggers the PUT
    await userEvent.click(screen.getByRole('button', { name: /overwrite/i }));
  };

  test('sends template_params when overwriting a dataset with include template parameters checked', async () => {
    // @ts-expect-error
    global.featureFlags = {
      [FeatureFlag.EnableTemplateProcessing]: true,
    };

    const putSpy = jest
      .spyOn(SupersetClient, 'put')
      .mockResolvedValue({ json: { result: { id: 0 } } } as any);

    const dummyDispatch = jest.fn().mockResolvedValue({});
    useDispatchMock.mockReturnValue(dummyDispatch);
    useSelectorMock.mockReturnValue({ ...user });

    const propsWithTemplateParam = {
      ...mockedProps,
      datasource: {
        ...testQuery,
        templateParams: JSON.stringify({ my_param: 12, _filters: 'foo' }),
      },
    };
    render(<SaveDatasetModal {...propsWithTemplateParam} />, {
      useRedux: true,
    });

    // Check the "Include Template Parameters" checkbox
    await userEvent.click(screen.getByRole('checkbox'));

    await setupOverwriteFlow();

    await waitFor(() => {
      expect(
        putSpy.mock.calls.some(([req]) =>
          req.endpoint?.includes('api/v1/dataset/'),
        ),
      ).toBe(true);
    });

    const datasetPutCall = putSpy.mock.calls.find(([req]) =>
      req.endpoint?.includes('api/v1/dataset/'),
    )!;
    const [req] = datasetPutCall;
    expect(req.endpoint).toContain('override_columns=true');
    const body = JSON.parse(req.body as string);
    // _filters should be stripped, but my_param should be preserved
    expect(body.template_params).toEqual(JSON.stringify({ my_param: 12 }));

    putSpy.mockRestore();
  });

  test('does not send template_params when overwriting a dataset with include template parameters unchecked', async () => {
    // @ts-expect-error
    global.featureFlags = {
      [FeatureFlag.EnableTemplateProcessing]: true,
    };

    const putSpy = jest
      .spyOn(SupersetClient, 'put')
      .mockResolvedValue({ json: { result: { id: 0 } } } as any);

    const dummyDispatch = jest.fn().mockResolvedValue({});
    useDispatchMock.mockReturnValue(dummyDispatch);
    useSelectorMock.mockReturnValue({ ...user });

    const propsWithTemplateParam = {
      ...mockedProps,
      datasource: {
        ...testQuery,
        templateParams: JSON.stringify({ my_param: 12 }),
      },
    };
    render(<SaveDatasetModal {...propsWithTemplateParam} />, {
      useRedux: true,
    });

    // Do NOT check the "Include Template Parameters" checkbox
    await setupOverwriteFlow();

    await waitFor(() => {
      expect(
        putSpy.mock.calls.some(([req]) =>
          req.endpoint?.includes('api/v1/dataset/'),
        ),
      ).toBe(true);
    });

    const datasetPutCall = putSpy.mock.calls.find(([req]) =>
      req.endpoint?.includes('api/v1/dataset/'),
    )!;
    const [req] = datasetPutCall;
    const body = JSON.parse(req.body as string);
    expect(body.template_params).toBeUndefined();

    putSpy.mockRestore();
  });

  test('clears dataset cache when creating new dataset', async () => {
    const clearDatasetCache = jest.spyOn(
      require('src/utils/cachedSupersetGet'),
      'clearDatasetCache',
    );
    const postFormData = jest.spyOn(
      require('src/explore/exploreUtils/formData'),
      'postFormData',
    );

    const dummyDispatch = jest.fn().mockResolvedValue({ id: 123 });
    useDispatchMock.mockReturnValue(dummyDispatch);
    useSelectorMock.mockReturnValue({ ...user });
    postFormData.mockResolvedValue('chart_key_123');

    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });

    const inputFieldText = screen.getByDisplayValue(/unimportant/i);
    fireEvent.change(inputFieldText, { target: { value: 'my dataset' } });

    const saveConfirmationBtn = screen.getByRole('button', {
      name: /save/i,
    });
    userEvent.click(saveConfirmationBtn);

    await waitFor(() => {
      expect(clearDatasetCache).toHaveBeenCalledWith(123);
    });
  });

  test('clearDatasetCache is imported and available', () => {
    const { clearDatasetCache } = require('src/utils/cachedSupersetGet');

    expect(clearDatasetCache).toBeDefined();
    expect(typeof clearDatasetCache).toBe('function');
  });
});