From 0fa66773e4f69fd50e977c0608738d22d99ac15e Mon Sep 17 00:00:00 2001
From: Diego Medina <diegomedina24@gmail.com>
Date: Fri, 25 Mar 2022 19:03:51 -0300
Subject: [PATCH] Add tests

---
 .../TableSelector/TableSelector.test.tsx      | 35 ++++++++++++++-----
 .../src/components/TableSelector/index.tsx    | 12 +++++--
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/superset-frontend/src/components/TableSelector/TableSelector.test.tsx b/superset-frontend/src/components/TableSelector/TableSelector.test.tsx
index ff8dc126d4064..32d84c008605c 100644
--- a/superset-frontend/src/components/TableSelector/TableSelector.test.tsx
+++ b/superset-frontend/src/components/TableSelector/TableSelector.test.tsx
@@ -22,7 +22,7 @@ import { render, screen, waitFor, within } from 'spec/helpers/testing-library';
 import { SupersetClient } from '@superset-ui/core';
 import { act } from 'react-dom/test-utils';
 import userEvent from '@testing-library/user-event';
-import TableSelector from '.';
+import TableSelector, { TableSelectorMultiple } from '.';
 
 const SupersetClientGet = jest.spyOn(SupersetClient, 'get');
 
@@ -55,6 +55,8 @@ const getTableMockFunction = async () =>
       options: [
         { label: 'table_a', value: 'table_a' },
         { label: 'table_b', value: 'table_b' },
+        { label: 'table_c', value: 'table_c' },
+        { label: 'table_d', value: 'table_d' },
       ],
     },
   } as any);
@@ -150,6 +152,8 @@ test('table options are notified after schema selection', async () => {
     expect(callback).toHaveBeenCalledWith([
       { label: 'table_a', value: 'table_a' },
       { label: 'table_b', value: 'table_b' },
+      { label: 'table_c', value: 'table_c' },
+      { label: 'table_d', value: 'table_d' },
     ]);
   });
 });
@@ -194,17 +198,15 @@ test('table select retain value if not in SQL Lab mode', async () => {
   ).toBeInTheDocument();
 });
 
-test('table select does not retain value in SQL Lab mode', async () => {
+test('table multi select retain all the values selected', async () => {
   SupersetClientGet.mockImplementation(getTableMockFunction);
-  document.body.innerHTML = '';
 
   const callback = jest.fn();
   const props = createProps({
     onTableSelectChange: callback,
-    sqlLabMode: true,
   });
 
-  render(<TableSelector {...props} />, { useRedux: true });
+  render(<TableSelectorMultiple {...props} />, { useRedux: true });
 
   const tableSelect = screen.getByRole('combobox', {
     name: 'Select table or type table name',
@@ -220,9 +222,26 @@ test('table select does not retain value in SQL Lab mode', async () => {
   ).toBeInTheDocument();
 
   act(() => {
-    userEvent.click(screen.getAllByText('table_a')[1]);
+    const item = screen.getAllByText('table_a');
+    userEvent.click(item[item.length - 1]);
   });
 
-  expect(callback).toHaveBeenCalled();
-  expect(getSelectItemContainer(tableSelect)).toHaveLength(0);
+  act(() => {
+    const item = screen.getAllByText('table_c');
+    userEvent.click(item[item.length - 1]);
+  });
+
+  const selectedValueContainer = getSelectItemContainer(tableSelect);
+
+  expect(selectedValueContainer).toHaveLength(2);
+  expect(
+    await within(selectedValueContainer?.[0] as HTMLElement).findByText(
+      'table_a',
+    ),
+  ).toBeInTheDocument();
+  expect(
+    await within(selectedValueContainer?.[1] as HTMLElement).findByText(
+      'table_c',
+    ),
+  ).toBeInTheDocument();
 });
diff --git a/superset-frontend/src/components/TableSelector/index.tsx b/superset-frontend/src/components/TableSelector/index.tsx
index aea1fc25ae494..c77d69e925311 100644
--- a/superset-frontend/src/components/TableSelector/index.tsx
+++ b/superset-frontend/src/components/TableSelector/index.tsx
@@ -23,6 +23,7 @@ import React, {
   useMemo,
   useEffect,
   useCallback,
+  useRef,
 } from 'react';
 import { SelectValue } from 'antd/lib/select';
 
@@ -94,7 +95,7 @@ interface TableSelectorProps {
   readOnly?: boolean;
   schema?: string;
   sqlLabMode?: boolean;
-  tableValue: string | string[];
+  tableValue?: string | string[];
   onTableSelectChange?: (value?: SelectValue, schema?: string) => void;
   tableSelectMode?: 'single' | 'multiple';
   emptyTableSelectValue?: SelectValue;
@@ -161,9 +162,11 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
   sqlLabMode = true,
   tableSelectMode = 'single',
   emptyTableSelectValue = undefined,
-  tableValue,
+  tableValue = undefined,
   onTableSelectChange,
 }) => {
+  const selectRef = useRef<HTMLInputElement>(null);
+
   const [currentDatabase, setCurrentDatabase] = useState<
     DatabaseObject | undefined
   >(database);
@@ -254,6 +257,10 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
   }
 
   const internalTableChange = (value: SelectValue | undefined) => {
+    if (tableSelectMode === 'multiple') {
+      selectRef?.current?.blur();
+    }
+
     if (currentSchema) {
       onTableSelectChange?.(value, currentSchema);
     } else {
@@ -318,6 +325,7 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
 
     const select = (
       <Select
+        ref={selectRef}
         ariaLabel={t('Select table or type table name')}
         disabled={disabled}
         filterOption={handleFilterOption}