[Data Table] Add unit tests (#90173)

* Move formatting columns into response handler

* Use shared csv export

* Cleanup files

* Fix type

* Fix translation

* Filter out non-dimension values

* Add unit tests for tableVisResponseHandler

* Add unit tests for createFormattedTable

* Add unit tests for addPercentageColumn

* Add unit tests for usePagination

* Add unit tests for useUiState

* Add unit tests for table visualization

* Add unit tests for TableVisBasic

* Add unit tests for cell

* Update license

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

src/plugins/vis_type_table/public/components/__snapshots__/table_vis_basic.test.tsx.snap

@@ -0,0 +1,115 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`TableVisBasic should init data grid 1`] = `
+<Fragment>
+  <EuiDataGrid
+    aria-label="Data table visualization"
+    columnVisibility={
+      Object {
+        "setVisibleColumns": [Function],
+        "visibleColumns": Array [],
+      }
+    }
+    columns={Array []}
+    gridStyle={
+      Object {
+        "border": "horizontal",
+        "header": "underline",
+      }
+    }
+    minSizeForControls={1}
+    onColumnResize={[Function]}
+    renderCellValue={[Function]}
+    rowCount={0}
+    sorting={
+      Object {
+        "columns": Array [],
+        "onSort": [Function],
+      }
+    }
+  />
+</Fragment>
+`;
+
+exports[`TableVisBasic should init data grid with title provided - for split mode 1`] = `
+<Fragment>
+  <EuiTitle
+    size="xs"
+  >
+    <h3>
+      My data table
+    </h3>
+  </EuiTitle>
+  <EuiDataGrid
+    aria-label="My data table"
+    columnVisibility={
+      Object {
+        "setVisibleColumns": [Function],
+        "visibleColumns": Array [],
+      }
+    }
+    columns={Array []}
+    gridStyle={
+      Object {
+        "border": "horizontal",
+        "header": "underline",
+      }
+    }
+    minSizeForControls={1}
+    onColumnResize={[Function]}
+    renderCellValue={[Function]}
+    rowCount={0}
+    sorting={
+      Object {
+        "columns": Array [],
+        "onSort": [Function],
+      }
+    }
+  />
+</Fragment>
+`;
+
+exports[`TableVisBasic should render the toolbar 1`] = `
+<Fragment>
+  <EuiDataGrid
+    aria-label="My saved vis"
+    columnVisibility={
+      Object {
+        "setVisibleColumns": [Function],
+        "visibleColumns": Array [],
+      }
+    }
+    columns={Array []}
+    gridStyle={
+      Object {
+        "border": "horizontal",
+        "header": "underline",
+      }
+    }
+    minSizeForControls={1}
+    onColumnResize={[Function]}
+    renderCellValue={[Function]}
+    rowCount={0}
+    sorting={
+      Object {
+        "columns": Array [],
+        "onSort": [Function],
+      }
+    }
+    toolbarVisibility={
+      Object {
+        "additionalControls": <Memo
+          columns={Array []}
+          dataGridAriaLabel="My saved vis"
+          filename="My saved vis"
+          rows={Array []}
+        />,
+        "showColumnSelector": false,
+        "showFullScreenSelector": false,
+        "showSortSelector": false,
+        "showStyleSelector": false,
+      }
+    }
+  />
+</Fragment>
+`;

src/plugins/vis_type_table/public/components/__snapshots__/table_vis_cell.test.tsx.snap

@@ -0,0 +1,13 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`table vis cell should return a cell component with data in scope 1`] = `
+<div
+  className="tbvChartCellContent"
+  dangerouslySetInnerHTML={
+    Object {
+      "__html": undefined,
+    }
+  }
+  data-test-subj="tbvChartCellContent"
+/>
+`;

src/plugins/vis_type_table/public/components/table_vis_basic.test.tsx

@@ -0,0 +1,130 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import React from 'react';
+import { shallow } from 'enzyme';
+import { TableVisBasic } from './table_vis_basic';
+import { FormattedColumn, TableVisConfig, TableVisUiState } from '../types';
+import { DatatableColumn } from 'src/plugins/expressions';
+import { createTableVisCell } from './table_vis_cell';
+import { createGridColumns } from './table_vis_columns';
+
+jest.mock('./table_vis_columns', () => ({
+  createGridColumns: jest.fn(() => []),
+}));
+jest.mock('./table_vis_cell', () => ({
+  createTableVisCell: jest.fn(() => () => {}),
+}));
+
+describe('TableVisBasic', () => {
+  const props = {
+    fireEvent: jest.fn(),
+    table: {
+      columns: [],
+      rows: [],
+      formattedColumns: {
+        test: {
+          formattedTotal: 100,
+        } as FormattedColumn,
+      },
+    },
+    visConfig: {} as TableVisConfig,
+    uiStateProps: {
+      sort: {
+        columnIndex: null,
+        direction: null,
+      },
+      columnsWidth: [],
+      setColumnsWidth: jest.fn(),
+      setSort: jest.fn(),
+    },
+  };
+
+  it('should init data grid', () => {
+    const comp = shallow(<TableVisBasic {...props} />);
+    expect(comp).toMatchSnapshot();
+  });
+
+  it('should init data grid with title provided - for split mode', () => {
+    const title = 'My data table';
+    const comp = shallow(<TableVisBasic {...props} title={title} />);
+    expect(comp).toMatchSnapshot();
+  });
+
+  it('should render the toolbar', () => {
+    const comp = shallow(
+      <TableVisBasic
+        {...props}
+        visConfig={{ ...props.visConfig, title: 'My saved vis', showToolbar: true }}
+      />
+    );
+    expect(comp).toMatchSnapshot();
+  });
+
+  it('should sort rows by column and pass the sorted rows for consumers', () => {
+    const uiStateProps = {
+      ...props.uiStateProps,
+      sort: {
+        columnIndex: 1,
+        direction: 'desc',
+      } as TableVisUiState['sort'],
+    };
+    const table = {
+      columns: [{ id: 'first' }, { id: 'second' }] as DatatableColumn[],
+      rows: [
+        { first: 1, second: 2 },
+        { first: 3, second: 4 },
+        { first: 5, second: 6 },
+      ],
+      formattedColumns: {},
+    };
+    const sortedRows = [
+      { first: 5, second: 6 },
+      { first: 3, second: 4 },
+      { first: 1, second: 2 },
+    ];
+    const comp = shallow(
+      <TableVisBasic
+        {...props}
+        table={table}
+        uiStateProps={uiStateProps}
+        visConfig={{ ...props.visConfig, showToolbar: true }}
+      />
+    );
+    expect(createTableVisCell).toHaveBeenCalledWith(sortedRows, table.formattedColumns);
+    expect(createGridColumns).toHaveBeenCalledWith(
+      table.columns,
+      sortedRows,
+      table.formattedColumns,
+      uiStateProps.columnsWidth,
+      props.fireEvent
+    );
+
+    const { onSort } = comp.find('EuiDataGrid').prop('sorting');
+    // sort the first column
+    onSort([{ id: 'first', direction: 'asc' }]);
+    expect(uiStateProps.setSort).toHaveBeenCalledWith({ columnIndex: 0, direction: 'asc' });
+    // sort the second column - should erase the first column sorting since there is only one level sorting available
+    onSort([
+      { id: 'first', direction: 'asc' },
+      { id: 'second', direction: 'desc' },
+    ]);
+    expect(uiStateProps.setSort).toHaveBeenCalledWith({ columnIndex: 1, direction: 'desc' });
+  });
+
+  it('should pass renderFooterCellValue for the total row', () => {
+    const comp = shallow(
+      <TableVisBasic {...props} visConfig={{ ...props.visConfig, showTotal: true }} />
+    );
+    const renderFooterCellValue: (props: any) => void = comp
+      .find('EuiDataGrid')
+      .prop('renderFooterCellValue');
+    expect(renderFooterCellValue).toEqual(expect.any(Function));
+    expect(renderFooterCellValue({ columnId: 'test' })).toEqual(100);
+  });
+});

src/plugins/vis_type_table/public/components/table_vis_cell.test.tsx

@@ -0,0 +1,36 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import React from 'react';
+import { shallow } from 'enzyme';
+import { EuiDataGridCellValueElementProps } from '@elastic/eui';
+import { createTableVisCell } from './table_vis_cell';
+import { FormattedColumns } from '../types';
+
+describe('table vis cell', () => {
+  it('should return a cell component with data in scope', () => {
+    const rows = [{ first: 1, second: 2 }];
+    const formattedColumns = ({
+      second: {
+        formatter: {
+          convert: jest.fn(),
+        },
+      },
+    } as unknown) as FormattedColumns;
+    const Cell = createTableVisCell(rows, formattedColumns);
+    const cellProps = {
+      rowIndex: 0,
+      columnId: 'second',
+    } as EuiDataGridCellValueElementProps;
+
+    const comp = shallow(<Cell {...cellProps} />);
+
+    expect(comp).toMatchSnapshot();
+    expect(formattedColumns.second.formatter.convert).toHaveBeenLastCalledWith(2, 'html');
+  });
+});

src/plugins/vis_type_table/public/components/table_visualization.test.tsx

@@ -0,0 +1,69 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+jest.mock('../utils', () => ({
+  useUiState: jest.fn(() => 'uiState'),
+}));
+
+import React from 'react';
+import { shallow } from 'enzyme';
+import { IInterpreterRenderHandlers } from 'src/plugins/expressions';
+import { coreMock } from '../../../../core/public/mocks';
+import { TableVisConfig, TableVisData } from '../types';
+import TableVisualizationComponent from './table_visualization';
+import { useUiState } from '../utils';
+
+describe('TableVisualizationComponent', () => {
+  const coreStartMock = coreMock.createStart();
+  const handlers = ({
+    done: jest.fn(),
+    uiState: 'uiState',
+    event: 'event',
+  } as unknown) as IInterpreterRenderHandlers;
+  const visData: TableVisData = {
+    table: {
+      columns: [],
+      rows: [],
+      formattedColumns: {},
+    },
+    tables: [],
+  };
+  const visConfig = ({} as unknown) as TableVisConfig;
+
+  it('should render the basic table', () => {
+    const comp = shallow(
+      <TableVisualizationComponent
+        core={coreStartMock}
+        handlers={handlers}
+        visData={visData}
+        visConfig={visConfig}
+      />
+    );
+    expect(useUiState).toHaveBeenLastCalledWith(handlers.uiState);
+    expect(comp.find('.tbvChart__splitColumns').exists()).toBeFalsy();
+    expect(comp.find('.tbvChart__split').exists()).toBeTruthy();
+  });
+
+  it('should render split table', () => {
+    const comp = shallow(
+      <TableVisualizationComponent
+        core={coreStartMock}
+        handlers={handlers}
+        visData={{
+          direction: 'column',
+          tables: [],
+        }}
+        visConfig={visConfig}
+      />
+    );
+    expect(useUiState).toHaveBeenLastCalledWith(handlers.uiState);
+    expect(comp.find('.tbvChart__splitColumns').exists()).toBeTruthy();
+    expect(comp.find('.tbvChart__split').exists()).toBeFalsy();
+    expect(comp.find('[data-test-subj="tbvChart"]').children().prop('tables')).toEqual([]);
+  });
+});

src/plugins/vis_type_table/public/utils/add_percentage_column.test.ts

@@ -0,0 +1,79 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+jest.mock('../services', () => ({
+  getFormatService: jest.fn(() => ({
+    deserialize: jest.fn(() => 'formatter'),
+  })),
+}));
+
+import { FieldFormat } from 'src/plugins/data/public';
+import { TableContext } from '../types';
+import { addPercentageColumn } from './add_percentage_column';
+
+describe('', () => {
+  const table: TableContext = {
+    columns: [
+      { id: 'col-0-1', name: 'Count', meta: { type: 'number' } },
+      { id: 'col-1-5', name: 'category.keyword: Descending', meta: { type: 'string' } },
+      { id: 'col-1-2', name: 'Gender', meta: { type: 'string' } },
+    ],
+    rows: [
+      { 'col-0-1': 1, 'col-1-5': "Women's Clothing", 'col-1-2': 'Men' },
+      { 'col-0-1': 6, 'col-1-5': "Women's Clothing", 'col-1-2': 'Men' },
+    ],
+    formattedColumns: {
+      'col-0-1': {
+        sumTotal: 7,
+        title: 'Count',
+        filterable: false,
+        formatter: {} as FieldFormat,
+      },
+    },
+  };
+
+  it('should dnot add percentage column if it was not found', () => {
+    const output = addPercentageColumn(table, 'Extra');
+    expect(output).toBe(table);
+  });
+
+  it('should add a brand new percentage column into table based on data', () => {
+    const output = addPercentageColumn(table, 'Count');
+    const expectedColumns = [
+      table.columns[0],
+      {
+        id: 'col-0-1-percents',
+        meta: {
+          params: {
+            id: 'percent',
+          },
+          type: 'number',
+        },
+        name: 'Count percentages',
+      },
+      table.columns[1],
+      table.columns[2],
+    ];
+    const expectedRows = [
+      { ...table.rows[0], 'col-0-1-percents': 0.14285714285714285 },
+      { ...table.rows[1], 'col-0-1-percents': 0.8571428571428571 },
+    ];
+    expect(output).toEqual({
+      columns: expectedColumns,
+      rows: expectedRows,
+      formattedColumns: {
+        ...table.formattedColumns,
+        'col-0-1-percents': {
+          filterable: false,
+          formatter: 'formatter',
+          title: 'Count percentages',
+        },
+      },
+    });
+  });
+});

src/plugins/vis_type_table/public/utils/create_formatted_table.test.ts

@@ -0,0 +1,218 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+const mockDeserialize = jest.fn(() => ({}));
+
+jest.mock('../services', () => ({
+  getFormatService: jest.fn(() => ({
+    deserialize: mockDeserialize,
+  })),
+}));
+
+import { Datatable } from 'src/plugins/expressions';
+import { AggTypes } from '../../common';
+import { TableVisConfig } from '../types';
+import { createFormattedTable } from './create_formatted_table';
+
+const visConfig: TableVisConfig = {
+  perPage: 10,
+  showPartialRows: false,
+  showMetricsAtAllLevels: false,
+  showToolbar: false,
+  showTotal: false,
+  totalFunc: AggTypes.SUM,
+  percentageCol: '',
+  title: 'My data table',
+  dimensions: {
+    buckets: [
+      {
+        accessor: 1,
+        aggType: 'terms',
+        format: { id: 'string' },
+        label: 'category_keyword: Descending',
+        params: {},
+      },
+    ],
+    metrics: [
+      { accessor: 0, aggType: 'count', format: { id: 'number' }, label: 'Count', params: {} },
+    ],
+  },
+};
+
+describe('createFormattedTable', () => {
+  const table: Datatable = {
+    columns: [
+      { id: 'col-0-1', name: 'Count', meta: { type: 'number' } },
+      { id: 'col-1-5', name: 'category.keyword: Descending', meta: { type: 'string' } },
+      { id: 'col-1-2', name: 'Gender', meta: { type: 'string' } },
+    ],
+    rows: [
+      { 'col-0-1': 1, 'col-1-5': "Women's Clothing", 'col-1-2': 'Men' },
+      { 'col-0-1': 6, 'col-1-5': "Women's Clothing", 'col-1-2': 'Men' },
+    ],
+    type: 'datatable',
+  };
+
+  it('should create formatted columns from data response and flter out non-dimension columns', () => {
+    const output = createFormattedTable(table, visConfig);
+
+    // column to split is filtered out of real data representing
+    expect(output.columns).toEqual([table.columns[0], table.columns[1]]);
+    expect(output.rows).toEqual(table.rows);
+    expect(output.formattedColumns).toEqual({
+      'col-0-1': {
+        filterable: false,
+        formatter: {},
+        title: 'Count',
+      },
+      'col-1-5': {
+        filterable: true,
+        formatter: {},
+        title: 'category.keyword: Descending',
+      },
+    });
+  });
+
+  it('should add total sum to numeric columns', () => {
+    mockDeserialize.mockImplementationOnce(() => ({
+      allowsNumericalAggregations: true,
+      convert: jest.fn((number) => number),
+    }));
+    const output = createFormattedTable(table, visConfig);
+
+    expect(output.formattedColumns).toEqual({
+      'col-0-1': {
+        filterable: false,
+        formatter: {
+          allowsNumericalAggregations: true,
+          convert: expect.any(Function),
+        },
+        title: 'Count',
+        sumTotal: 7,
+        total: 7,
+        formattedTotal: 7,
+      },
+      'col-1-5': {
+        filterable: true,
+        formatter: {},
+        title: 'category.keyword: Descending',
+      },
+    });
+  });
+
+  it('should add total average to numeric columns', () => {
+    mockDeserialize.mockImplementationOnce(() => ({
+      allowsNumericalAggregations: true,
+      convert: jest.fn((number) => number),
+    }));
+    const output = createFormattedTable(table, { ...visConfig, totalFunc: AggTypes.AVG });
+
+    expect(output.formattedColumns).toEqual({
+      'col-0-1': {
+        filterable: false,
+        formatter: {
+          allowsNumericalAggregations: true,
+          convert: expect.any(Function),
+        },
+        title: 'Count',
+        sumTotal: 7,
+        total: 3.5,
+        formattedTotal: 3.5,
+      },
+      'col-1-5': {
+        filterable: true,
+        formatter: {},
+        title: 'category.keyword: Descending',
+      },
+    });
+  });
+
+  it('should find min value as total', () => {
+    mockDeserialize.mockImplementationOnce(() => ({
+      allowsNumericalAggregations: true,
+      convert: jest.fn((number) => number),
+    }));
+    const output = createFormattedTable(table, { ...visConfig, totalFunc: AggTypes.MIN });
+
+    expect(output.formattedColumns).toEqual({
+      'col-0-1': {
+        filterable: false,
+        formatter: {
+          allowsNumericalAggregations: true,
+          convert: expect.any(Function),
+        },
+        title: 'Count',
+        sumTotal: 7,
+        total: 1,
+        formattedTotal: 1,
+      },
+      'col-1-5': {
+        filterable: true,
+        formatter: {},
+        title: 'category.keyword: Descending',
+      },
+    });
+  });
+
+  it('should find max value as total', () => {
+    mockDeserialize.mockImplementationOnce(() => ({
+      allowsNumericalAggregations: true,
+      convert: jest.fn((number) => number),
+    }));
+    const output = createFormattedTable(table, { ...visConfig, totalFunc: AggTypes.MAX });
+
+    expect(output.formattedColumns).toEqual({
+      'col-0-1': {
+        filterable: false,
+        formatter: {
+          allowsNumericalAggregations: true,
+          convert: expect.any(Function),
+        },
+        title: 'Count',
+        sumTotal: 7,
+        total: 6,
+        formattedTotal: 6,
+      },
+      'col-1-5': {
+        filterable: true,
+        formatter: {},
+        title: 'category.keyword: Descending',
+      },
+    });
+  });
+
+  it('should add rows count as total', () => {
+    mockDeserialize.mockImplementationOnce(() => ({
+      allowsNumericalAggregations: true,
+      convert: jest.fn((number) => number),
+    }));
+    const output = createFormattedTable(table, { ...visConfig, totalFunc: AggTypes.COUNT });
+
+    expect(output.formattedColumns).toEqual({
+      'col-0-1': {
+        filterable: false,
+        formatter: {
+          allowsNumericalAggregations: true,
+          convert: expect.any(Function),
+        },
+        title: 'Count',
+        sumTotal: 7,
+        total: 2,
+        formattedTotal: 2,
+      },
+      'col-1-5': {
+        filterable: true,
+        formattedTotal: 2,
+        formatter: {},
+        sumTotal: "0Women's ClothingWomen's Clothing",
+        title: 'category.keyword: Descending',
+        total: 2,
+      },
+    });
+  });
+});

src/plugins/vis_type_table/public/utils/table_vis_response_handler.test.ts

@@ -0,0 +1,171 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+const mockConverter = jest.fn((name) => `By ${name}`);
+
+jest.mock('../services', () => ({
+  getFormatService: jest.fn(() => ({
+    deserialize: jest.fn(() => ({
+      convert: mockConverter,
+    })),
+  })),
+}));
+
+jest.mock('./create_formatted_table', () => ({
+  createFormattedTable: jest.fn((data) => ({
+    ...data,
+    formattedColumns: {},
+  })),
+}));
+
+jest.mock('./add_percentage_column', () => ({
+  addPercentageColumn: jest.fn((data, column) => ({
+    ...data,
+    percentage: `${column} with percentage`,
+  })),
+}));
+
+import { Datatable } from 'src/plugins/expressions';
+import { SchemaConfig } from 'src/plugins/visualizations/public';
+import { AggTypes } from '../../common';
+import { TableGroup, TableVisConfig } from '../types';
+import { addPercentageColumn } from './add_percentage_column';
+import { createFormattedTable } from './create_formatted_table';
+import { tableVisResponseHandler } from './table_vis_response_handler';
+
+const visConfig: TableVisConfig = {
+  perPage: 10,
+  showPartialRows: false,
+  showMetricsAtAllLevels: false,
+  showToolbar: false,
+  showTotal: false,
+  totalFunc: AggTypes.AVG,
+  percentageCol: '',
+  title: 'My data table',
+  dimensions: {
+    buckets: [],
+    metrics: [],
+  },
+};
+
+describe('tableVisResponseHandler', () => {
+  describe('basic table', () => {
+    const input: Datatable = {
+      columns: [],
+      rows: [],
+      type: 'datatable',
+    };
+
+    it('should create formatted table for basic usage', () => {
+      const output = tableVisResponseHandler(input, visConfig);
+
+      expect(output.direction).toBeUndefined();
+      expect(output.tables.length).toEqual(0);
+      expect(addPercentageColumn).not.toHaveBeenCalled();
+      expect(createFormattedTable).toHaveBeenCalledWith(input, visConfig);
+      expect(output.table).toEqual({
+        ...input,
+        formattedColumns: {},
+      });
+    });
+
+    it('should add a percentage column if it is set', () => {
+      const output = tableVisResponseHandler(input, { ...visConfig, percentageCol: 'Count' });
+      expect(output.table).toEqual({
+        ...input,
+        formattedColumns: {},
+        percentage: 'Count with percentage',
+      });
+    });
+  });
+
+  describe('split table', () => {
+    const input: Datatable = {
+      columns: [
+        { id: 'col-0-1', name: 'Count', meta: { type: 'number' } },
+        { id: 'col-1-2', name: 'Gender', meta: { type: 'string' } },
+      ],
+      rows: [
+        { 'col-0-1': 1, 'col-1-2': 'Men' },
+        { 'col-0-1': 3, 'col-1-2': 'Women' },
+        { 'col-0-1': 6, 'col-1-2': 'Men' },
+      ],
+      type: 'datatable',
+    };
+    const split: SchemaConfig[] = [
+      {
+        accessor: 1,
+        label: 'Split',
+        format: {},
+        params: {},
+        aggType: 'terms',
+      },
+    ];
+    const expectedOutput: TableGroup[] = [
+      {
+        title: 'By Men: Gender',
+        table: {
+          columns: input.columns,
+          rows: [input.rows[0], input.rows[2]],
+          formattedColumns: {},
+        },
+      },
+      {
+        title: 'By Women: Gender',
+        table: {
+          columns: input.columns,
+          rows: [input.rows[1]],
+          formattedColumns: {},
+        },
+      },
+    ];
+
+    it('should split data by row', () => {
+      const output = tableVisResponseHandler(input, {
+        ...visConfig,
+        dimensions: { ...visConfig.dimensions, splitRow: split },
+      });
+
+      expect(output.direction).toEqual('row');
+      expect(output.table).toBeUndefined();
+      expect(output.tables).toEqual(expectedOutput);
+    });
+
+    it('should split data by column', () => {
+      const output = tableVisResponseHandler(input, {
+        ...visConfig,
+        dimensions: { ...visConfig.dimensions, splitColumn: split },
+      });
+
+      expect(output.direction).toEqual('column');
+      expect(output.table).toBeUndefined();
+      expect(output.tables).toEqual(expectedOutput);
+    });
+
+    it('should add  percentage columns to each table', () => {
+      const output = tableVisResponseHandler(input, {
+        ...visConfig,
+        percentageCol: 'Count',
+        dimensions: { ...visConfig.dimensions, splitColumn: split },
+      });
+
+      expect(output.direction).toEqual('column');
+      expect(output.table).toBeUndefined();
+      expect(output.tables).toEqual([
+        {
+          ...expectedOutput[0],
+          table: { ...expectedOutput[0].table, percentage: 'Count with percentage' },
+        },
+        {
+          ...expectedOutput[1],
+          table: { ...expectedOutput[1].table, percentage: 'Count with percentage' },
+        },
+      ]);
+    });
+  });
+});

src/plugins/vis_type_table/public/utils/table_vis_response_handler.ts

@@ -27,7 +27,6 @@ export function tableVisResponseHandler(input: Datatable, visConfig: TableVisCon
     const splitColumnIndex = split[0].accessor;
     const splitColumnFormatter = getFormatService().deserialize(split[0].format);
     const splitColumn = input.columns[splitColumnIndex];
-    const columns = input.columns.filter((c, idx) => idx !== splitColumnIndex);
     const splitMap: { [key: string]: number } = {};
     let splitIndex = 0;
 
@@ -39,7 +38,7 @@ export function tableVisResponseHandler(input: Datatable, visConfig: TableVisCon
         const tableGroup: TableGroup = {
           title: `${splitColumnFormatter.convert(splitValue)}: ${splitColumn.name}`,
           table: {
-            columns,
+            columns: input.columns,
             rows: [],
             formattedColumns: {},
           },
@@ -53,7 +52,7 @@ export function tableVisResponseHandler(input: Datatable, visConfig: TableVisCon
     });
 
     tables.forEach((tg) => {
-      tg.table = createFormattedTable({ ...tg.table, columns: input.columns }, visConfig);
+      tg.table = createFormattedTable(tg.table, visConfig);
 
       if (visConfig.percentageCol) {
         tg.table = addPercentageColumn(tg.table, visConfig.percentageCol);

src/plugins/vis_type_table/public/utils/use/use_pagination.test.ts

@@ -0,0 +1,119 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { act, renderHook } from '@testing-library/react-hooks';
+import { AggTypes } from '../../../common';
+import { usePagination } from './use_pagination';
+
+describe('usePagination', () => {
+  const visParams = {
+    perPage: 10,
+    showPartialRows: false,
+    showMetricsAtAllLevels: false,
+    showToolbar: false,
+    showTotal: false,
+    totalFunc: AggTypes.SUM,
+    percentageCol: '',
+    title: 'My data table',
+  };
+
+  it('should set up pagination on init', () => {
+    const { result } = renderHook(() => usePagination(visParams, 15));
+
+    expect(result.current).toEqual({
+      pageIndex: 0,
+      pageSize: 10,
+      onChangeItemsPerPage: expect.any(Function),
+      onChangePage: expect.any(Function),
+    });
+  });
+
+  it('should skip setting the pagination if perPage is not set', () => {
+    const { result } = renderHook(() => usePagination({ ...visParams, perPage: '' }, 15));
+
+    expect(result.current).toBeUndefined();
+  });
+
+  it('should change the page via callback', () => {
+    const { result } = renderHook(() => usePagination(visParams, 15));
+
+    act(() => {
+      // change the page to the next one
+      result.current?.onChangePage(1);
+    });
+
+    expect(result.current).toEqual({
+      pageIndex: 1,
+      pageSize: 10,
+      onChangeItemsPerPage: expect.any(Function),
+      onChangePage: expect.any(Function),
+    });
+  });
+
+  it('should change items per page via callback', () => {
+    const { result } = renderHook(() => usePagination(visParams, 15));
+
+    act(() => {
+      // change the page to the next one
+      result.current?.onChangeItemsPerPage(20);
+    });
+
+    expect(result.current).toEqual({
+      pageIndex: 0,
+      pageSize: 20,
+      onChangeItemsPerPage: expect.any(Function),
+      onChangePage: expect.any(Function),
+    });
+  });
+
+  it('should change the page when props were changed', () => {
+    const { result, rerender } = renderHook(
+      (props) => usePagination(props.visParams, props.rowCount),
+      {
+        initialProps: {
+          visParams,
+          rowCount: 15,
+        },
+      }
+    );
+    const updatedParams = { ...visParams, perPage: 5 };
+
+    // change items per page count
+    rerender({ visParams: updatedParams, rowCount: 15 });
+
+    expect(result.current).toEqual({
+      pageIndex: 0,
+      pageSize: 5,
+      onChangeItemsPerPage: expect.any(Function),
+      onChangePage: expect.any(Function),
+    });
+
+    act(() => {
+      // change the page to the last one - 3
+      result.current?.onChangePage(3);
+    });
+
+    expect(result.current).toEqual({
+      pageIndex: 3,
+      pageSize: 5,
+      onChangeItemsPerPage: expect.any(Function),
+      onChangePage: expect.any(Function),
+    });
+
+    // decrease the rows count
+    rerender({ visParams: updatedParams, rowCount: 10 });
+
+    // should switch to the last available page
+    expect(result.current).toEqual({
+      pageIndex: 1,
+      pageSize: 5,
+      onChangeItemsPerPage: expect.any(Function),
+      onChangePage: expect.any(Function),
+    });
+  });
+});

src/plugins/vis_type_table/public/utils/use/use_ui_state.test.ts

@@ -0,0 +1,163 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { act, renderHook } from '@testing-library/react-hooks';
+import type { PersistedState } from 'src/plugins/visualizations/public';
+import { TableVisUiState } from '../../types';
+import { useUiState } from './use_ui_state';
+
+describe('useUiState', () => {
+  let uiState: PersistedState;
+
+  beforeEach(() => {
+    uiState = {
+      get: jest.fn(),
+      on: jest.fn(),
+      off: jest.fn(),
+      set: jest.fn(),
+    } as any;
+  });
+
+  it("should init default columnsWidth & sort if uiState doesn't have it set", () => {
+    const { result } = renderHook(() => useUiState(uiState));
+
+    expect(result.current).toEqual({
+      columnsWidth: [],
+      sort: {
+        columnIndex: null,
+        direction: null,
+      },
+      setColumnsWidth: expect.any(Function),
+      setSort: expect.any(Function),
+    });
+  });
+
+  it('should subscribe on uiState changes and update local state', async () => {
+    const { result, unmount, waitForNextUpdate } = renderHook(() => useUiState(uiState));
+
+    expect(uiState.on).toHaveBeenCalledWith('change', expect.any(Function));
+    // @ts-expect-error
+    const updateOnChange = uiState.on.mock.calls[0][1];
+
+    uiState.getChanges = jest.fn(() => ({
+      vis: {
+        params: {
+          sort: {
+            columnIndex: 1,
+            direction: 'asc',
+          },
+          colWidth: [],
+        },
+      },
+    }));
+
+    act(() => {
+      updateOnChange();
+    });
+
+    await waitForNextUpdate();
+
+    // should update local state with new values
+    expect(result.current).toEqual({
+      columnsWidth: [],
+      sort: {
+        columnIndex: 1,
+        direction: 'asc',
+      },
+      setColumnsWidth: expect.any(Function),
+      setSort: expect.any(Function),
+    });
+
+    act(() => {
+      updateOnChange();
+    });
+
+    // should skip setting the state again if it is equal
+    expect(result.current).toEqual({
+      columnsWidth: [],
+      sort: {
+        columnIndex: 1,
+        direction: 'asc',
+      },
+      setColumnsWidth: expect.any(Function),
+      setSort: expect.any(Function),
+    });
+
+    unmount();
+
+    expect(uiState.off).toHaveBeenCalledWith('change', updateOnChange);
+  });
+
+  describe('updating uiState through callbacks', () => {
+    beforeAll(() => {
+      jest.useFakeTimers();
+    });
+
+    it('should update the uiState with new sort', async () => {
+      const { result } = renderHook(() => useUiState(uiState));
+      const newSort: TableVisUiState['sort'] = {
+        columnIndex: 5,
+        direction: 'desc',
+      };
+
+      act(() => {
+        result.current.setSort(newSort);
+      });
+
+      expect(result.current.sort).toEqual(newSort);
+
+      jest.runAllTimers();
+
+      expect(uiState.set).toHaveBeenCalledTimes(1);
+      expect(uiState.set).toHaveBeenCalledWith('vis.params.sort', newSort);
+    });
+
+    it('should update the uiState with new columns width', async () => {
+      const { result } = renderHook(() => useUiState(uiState));
+      const col1 = { colIndex: 0, width: 300 };
+      const col2 = { colIndex: 1, width: 100 };
+
+      // set width of a column
+      act(() => {
+        result.current.setColumnsWidth(col1);
+      });
+
+      expect(result.current.columnsWidth).toEqual([col1]);
+
+      jest.runAllTimers();
+
+      expect(uiState.set).toHaveBeenCalledTimes(1);
+      expect(uiState.set).toHaveBeenLastCalledWith('vis.params.colWidth', [col1]);
+
+      // set width of another column
+      act(() => {
+        result.current.setColumnsWidth(col2);
+      });
+
+      jest.runAllTimers();
+
+      expect(uiState.set).toHaveBeenCalledTimes(2);
+      expect(uiState.set).toHaveBeenLastCalledWith('vis.params.colWidth', [col1, col2]);
+
+      const updatedCol1 = { colIndex: 0, width: 200 };
+      // update width of existing column
+      act(() => {
+        result.current.setColumnsWidth(updatedCol1);
+      });
+
+      jest.runAllTimers();
+
+      expect(uiState.set).toHaveBeenCalledTimes(3);
+      expect(uiState.set).toHaveBeenCalledWith('vis.params.colWidth', [updatedCol1, col2]);
+    });
+
+    afterAll(() => {
+      jest.useRealTimers();
+    });
+  });
+});