diff --git a/packages/ui/__stories__/Table.stories.tsx b/packages/ui/__stories__/Table.stories.tsx
index ec1ca5556..802811788 100644
--- a/packages/ui/__stories__/Table.stories.tsx
+++ b/packages/ui/__stories__/Table.stories.tsx
@@ -9,7 +9,7 @@ import { TextField } from '../src'
 
 import storyStyles from './TextField.stories.module.scss'
 import styles from './utils.scss'
-import { SortingRule } from 'react-table'
+import { SortingRule } from '../lib/Table/types'
 
 export default {
   title: 'Components/Table',
@@ -164,7 +164,7 @@ export const WithControlledSorting = () => {
       // Only update the data if this is the latest fetch
       if (fetchId === fetchIdRef.current) {
         const { id, desc: isDescending } = sortBy[0]
-        const sortColumn = id as keyof Person
+        const sortColumn = id
         const newData = [...bigDataSet]
         newData.sort((a, b) =>
           isDescending ? Number(b[sortColumn]) - Number(a[sortColumn]) : Number(a[sortColumn]) - Number(b[sortColumn]),
diff --git a/packages/ui/__tests__/.eslintrc.js b/packages/ui/__tests__/.eslintrc.js
index 491bd2df4..3f610df7d 100644
--- a/packages/ui/__tests__/.eslintrc.js
+++ b/packages/ui/__tests__/.eslintrc.js
@@ -6,5 +6,13 @@ module.exports = {
     '@typescript-eslint/no-unsafe-return': 'off',
     '@typescript-eslint/no-non-null-assertion': 'off',
     '@typescript-eslint/no-empty-function': 'off',
+    '@typescript-eslint/no-unused-vars': [
+      'error',
+      {
+        caughtErrors: 'none',
+        argsIgnorePattern: '_',
+        ignoreRestSiblings: true,
+      },
+    ],
   },
 }
diff --git a/packages/ui/__tests__/Alert.test.tsx b/packages/ui/__tests__/Alert.test.tsx
index c47a8298d..baced569e 100644
--- a/packages/ui/__tests__/Alert.test.tsx
+++ b/packages/ui/__tests__/Alert.test.tsx
@@ -75,7 +75,7 @@ describe('Alert', () => {
       expect(AlertElement.findDataTest('alert-content').props()).toMatchObject({
         children: content,
       })
-      expect(wrapper.findDataTest('alert-icon').props()).toMatchObject({
+      expect(wrapper.findDataTestFirst('alert-icon').props()).toMatchObject({
         icon,
         ariaLabel,
       })
diff --git a/packages/ui/__tests__/Badge.test.tsx b/packages/ui/__tests__/Badge.test.tsx
index 27112420d..593b09807 100644
--- a/packages/ui/__tests__/Badge.test.tsx
+++ b/packages/ui/__tests__/Badge.test.tsx
@@ -12,7 +12,7 @@ describe('Badge', () => {
   it('Renders Badge with all necessary components', async () => {
     const wrapper = await mountAndCheckA11Y(<Badge type="neutral" size="normal" content={badgeContent} />)
 
-    expect(wrapper.findDataTest('badge-icon').prop('size')).toBe('medium')
+    expect(wrapper.findDataTestFirst('badge-icon').prop('size')).toBe('medium')
     expect(wrapper.findDataTest('badge-content').text()).toBe(badgeContent)
   })
 
@@ -63,7 +63,7 @@ describe('Badge', () => {
     const wrapper = await mountAndCheckA11Y(<Badge type={type} size="small" content={badgeContent} />)
 
     expect(wrapper.findDataTest('badge-container').prop('className')).toContain(className)
-    expect(wrapper.findDataTest('badge-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('badge-icon').props()).toMatchObject({
       ariaLabel,
       icon,
       size: 'small',
diff --git a/packages/ui/__tests__/Calendar/CalendarRange.test.tsx b/packages/ui/__tests__/Calendar/CalendarRange.test.tsx
index ccbbc7ee7..880ad3992 100644
--- a/packages/ui/__tests__/Calendar/CalendarRange.test.tsx
+++ b/packages/ui/__tests__/Calendar/CalendarRange.test.tsx
@@ -28,7 +28,7 @@ describe('CalendarRange', () => {
     })
 
     // Icon
-    expect(wrapper.findDataTest('calendar-range-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('calendar-range-icon').props()).toMatchObject({
       icon: ArrowRight,
       ariaLabel: 'Arrow Right',
     })
@@ -72,7 +72,7 @@ describe('CalendarRange', () => {
     })
 
     // Icon
-    expect(wrapper.findDataTest('calendar-range-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('calendar-range-icon').props()).toMatchObject({
       icon: ArrowRight,
       ariaLabel: 'Arrow Right',
     })
diff --git a/packages/ui/__tests__/Card.test.tsx b/packages/ui/__tests__/Card.test.tsx
index b2dc5b8b1..2e4e2f094 100644
--- a/packages/ui/__tests__/Card.test.tsx
+++ b/packages/ui/__tests__/Card.test.tsx
@@ -59,8 +59,7 @@ describe('Card', () => {
     expect(content).toBeInTheDocument()
     expect(content.className).toEqual(styles.content)
     expect(within(content).queryByText(cardContent)).toBeInTheDocument()
-
-    expect(screen.queryByTestId('card-heading-icon')).not.toBeInTheDocument()
+    expect(screen.queryByTestId('card-heading-icon')).toBeInTheDocument()
   })
 
   it('renders caption', async () => {
diff --git a/packages/ui/__tests__/EmptyState.test.tsx b/packages/ui/__tests__/EmptyState.test.tsx
index 3fff7b7f2..d99a2cdb5 100644
--- a/packages/ui/__tests__/EmptyState.test.tsx
+++ b/packages/ui/__tests__/EmptyState.test.tsx
@@ -31,7 +31,7 @@ describe('EmptyState', () => {
     )
 
     expect(wrapper.findDataTest('empty-state-container').prop('className')).toBe(cn(styles.container))
-    expect(wrapper.findDataTest('empty-state-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('empty-state-icon').props()).toMatchObject({
       icon,
       ariaLabel: iconLabel,
       size: 'large',
@@ -59,7 +59,7 @@ describe('EmptyState', () => {
     )
 
     expect(wrapper.findDataTest('empty-state-container').prop('className')).toBe(cn(styles.container))
-    expect(wrapper.findDataTest('empty-state-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('empty-state-icon').props()).toMatchObject({
       icon,
       ariaLabel: iconLabel,
       size: 'large',
@@ -80,7 +80,7 @@ describe('EmptyState', () => {
     )
 
     expect(wrapper.findDataTest('empty-state-container').prop('className')).toBe(cn(styles.container, styles.large))
-    expect(wrapper.findDataTest('empty-state-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('empty-state-icon').props()).toMatchObject({
       icon,
       ariaLabel: iconLabel,
       size: 'xlarge',
@@ -108,7 +108,7 @@ describe('EmptyState', () => {
     )
 
     expect(wrapper.findDataTest('empty-state-container').prop('className')).toBe(cn(styles.container, styles.horizontal))
-    expect(wrapper.findDataTest('empty-state-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('empty-state-icon').props()).toMatchObject({
       icon,
       ariaLabel: iconLabel,
       size: 'large',
diff --git a/packages/ui/__tests__/Modal.test.tsx b/packages/ui/__tests__/Modal.test.tsx
index f17c0ebc8..67e97d1a0 100644
--- a/packages/ui/__tests__/Modal.test.tsx
+++ b/packages/ui/__tests__/Modal.test.tsx
@@ -98,7 +98,7 @@ describe('Modal', () => {
       onRequestClose: onClose,
     })
 
-    expect(wrapper.findDataTest('modal-header-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('modal-header-icon').props()).toMatchObject({
       icon: CloudLightning,
       ariaLabel: iconAriaLabel,
     })
diff --git a/packages/ui/__tests__/Notification.test.tsx b/packages/ui/__tests__/Notification.test.tsx
index 2a9b9264e..4990fd295 100644
--- a/packages/ui/__tests__/Notification.test.tsx
+++ b/packages/ui/__tests__/Notification.test.tsx
@@ -51,7 +51,7 @@ describe('Notification', () => {
     expect(wrapper.findDataTest('notification').props()).toMatchObject({
       className: cn(styles.notification, className),
     })
-    expect(wrapper.findDataTest('notification-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('notification-icon').props()).toMatchObject({
       ariaLabel,
       icon,
     })
@@ -72,7 +72,7 @@ describe('Notification', () => {
     expect(wrapper.findDataTest('notification').props()).toMatchObject({
       className: cn(styles.notification, styles.success),
     })
-    expect(wrapper.findDataTest('notification-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('notification-icon').props()).toMatchObject({
       ariaLabel: 'Check circle icon',
       icon: CheckCircle,
     })
diff --git a/packages/ui/__tests__/Overlay.test.tsx b/packages/ui/__tests__/Overlay.test.tsx
index 06f968fae..5ff78e87a 100644
--- a/packages/ui/__tests__/Overlay.test.tsx
+++ b/packages/ui/__tests__/Overlay.test.tsx
@@ -32,7 +32,7 @@ describe('Overlay', () => {
     const reactModal = wrapper.find(ReactModal)
     const overlayWrapper = reactModal.findDataTest('overlay-wrapper')
     const header = overlayWrapper.findDataTest('overlay-header')
-    const headerIcon = header.findDataTest('overlay-header-icon')
+    const headerIcon = header.findDataTestFirst('overlay-header-icon')
     const headerTitle = header.findDataTest('overlay-header-title')
     const headerCancelButton = header.findDataTestFirst('overlay-header-cancel-button')
     const overlayContent = overlayWrapper.findDataTest('overlay-content')
diff --git a/packages/ui/__tests__/Pagination.test.tsx b/packages/ui/__tests__/Pagination.test.tsx
index 59d32a9c0..271cd72b0 100644
--- a/packages/ui/__tests__/Pagination.test.tsx
+++ b/packages/ui/__tests__/Pagination.test.tsx
@@ -94,6 +94,7 @@ const nextPrevButtonPropsBase: IconButtonProps = {
   icon: ChevronRight,
   ariaLabel: 'Base',
   className: styles.iconButton,
+  'data-test': 'pagination-buttons-next-page',
 
   onClick: expect.anything(),
 }
@@ -139,6 +140,7 @@ describe('Pagination', () => {
       label: 'Rows',
       options: pageSizeOptions.map((opt) => ({ value: opt, label: opt.toString() })),
       size: 'small',
+      'data-test': 'pagination-rows-per-page-select',
 
       onChange: expect.anything(),
     })
@@ -158,30 +160,48 @@ describe('Pagination', () => {
 
     const buttons = wrapper.findDataTest('pagination-buttons').find(Button)
     expect(buttons).toHaveLength(6)
-    expect(buttons.at(0).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       className: `${styles.button} ${styles.selected}`,
       children: '1',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(0).prop(key)).toEqual(value)
     })
-    expect(buttons.at(1).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '2',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(1).prop(key)).toEqual(value)
     })
-    expect(buttons.at(2).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '3',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(2).prop(key)).toEqual(value)
     })
-    expect(buttons.at(3).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '4',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(3).prop(key)).toEqual(value)
     })
-    expect(buttons.at(4).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '5',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(4).prop(key)).toEqual(value)
     })
-    expect(buttons.at(5).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '2000',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(5).prop(key)).toEqual(value)
     })
 
     expect(wrapper.find(NumberField).props()).toEqual<NumberFieldProps>({
@@ -194,7 +214,7 @@ describe('Pagination', () => {
       max: pageCount,
       size: 'small',
       value: currentPage,
-
+      'data-test': 'pagination-jump-from-input-field',
       onChange: expect.anything(),
     })
   })
@@ -230,39 +250,61 @@ describe('Pagination', () => {
 
     const iconButton = wrapper.findDataTest('pagination-buttons').find(IconButton)
     expect(iconButton).toHaveLength(2)
-    expect(iconButton.at(0).props()).toEqual<IconButtonProps>({
+
+    Object.entries({
       ...nextPrevButtonPropsBase,
       icon: ChevronLeft,
       ariaLabel: 'Previous page',
+      'data-test': 'pagination-buttons-previous-page',
+    }).forEach(([key, value]) => {
+      expect(iconButton.at(0).prop(key)).toEqual(value)
     })
-    expect(iconButton.at(1).props()).toEqual<IconButtonProps>({
+
+    Object.entries({
       ...nextPrevButtonPropsBase,
       icon: ChevronRight,
       ariaLabel: 'Next page',
+    }).forEach(([key, value]) => {
+      expect(iconButton.at(1).prop(key)).toEqual(value)
     })
 
     const buttons = wrapper.find(Button)
     expect(buttons).toHaveLength(5)
-    expect(buttons.at(0).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '1',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(0).prop(key)).toEqual(value)
     })
-    expect(buttons.at(1).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '999',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(1).prop(key)).toEqual(value)
     })
-    expect(buttons.at(2).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       className: `${styles.button} ${styles.selected}`,
       children: '1000',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(2).prop(key)).toEqual(value)
     })
-    expect(buttons.at(3).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '1001',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(3).prop(key)).toEqual(value)
     })
-    expect(buttons.at(4).props()).toEqual<ButtonProps>({
+    Object.entries({
       ...buttonPropsBase,
       children: '2000',
+      'data-test': 'pagination-buttons-go-to-page',
+    }).forEach(([key, value]) => {
+      expect(buttons.at(4).prop(key)).toEqual(value)
     })
   })
 
@@ -300,6 +342,7 @@ describe('Pagination', () => {
       ...nextPrevButtonPropsBase,
       icon: ChevronLeft,
       ariaLabel: 'Previous page',
+      'data-test': 'pagination-buttons-previous-page',
     })
 
     const buttons = wrapper.find(Button)
@@ -307,27 +350,33 @@ describe('Pagination', () => {
     expect(buttons.at(0).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '1',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(1).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '1996',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(2).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '1997',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(3).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '1998',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(4).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '1999',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(5).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       className: `${styles.button} ${styles.selected}`,
       children: '2000',
+      'data-test': 'pagination-buttons-go-to-page',
     })
   })
 
@@ -502,14 +551,17 @@ describe('Pagination', () => {
       ...buttonPropsBase,
       className: `${styles.button} ${styles.selected}`,
       children: '1',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(1).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '2',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(2).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '2000',
+      'data-test': 'pagination-buttons-go-to-page',
     })
 
     expect(wrapper.findDataTest('pagination-range-of-shown-items').exists()).toBe(false)
@@ -534,6 +586,7 @@ describe('Pagination', () => {
       label: 'Rows',
       size: 'small',
       options: pageSizeOptions.map((opt) => ({ value: opt, label: opt.toString() })),
+      'data-test': 'pagination-rows-per-page-select',
 
       onChange: expect.anything(),
     })
@@ -548,6 +601,7 @@ describe('Pagination', () => {
       max: pageCount,
       size: 'small',
       value: currentPage,
+      'data-test': 'pagination-jump-from-input-field',
 
       onChange: expect.anything(),
     })
@@ -570,14 +624,17 @@ describe('Pagination', () => {
       ...buttonPropsBase,
       className: `${styles.button} ${styles.selected}`,
       children: '1',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(1).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '2',
+      'data-test': 'pagination-buttons-go-to-page',
     })
     expect(buttons.at(2).props()).toEqual<ButtonProps>({
       ...buttonPropsBase,
       children: '2000',
+      'data-test': 'pagination-buttons-go-to-page',
     })
 
     expect(wrapper.findDataTest('pagination-range-of-shown-items').exists()).toBe(false)
diff --git a/packages/ui/__tests__/Select/SelectField.test.tsx b/packages/ui/__tests__/Select/SelectField.test.tsx
index 136aa4fab..179a21a1c 100644
--- a/packages/ui/__tests__/Select/SelectField.test.tsx
+++ b/packages/ui/__tests__/Select/SelectField.test.tsx
@@ -199,7 +199,7 @@ describe('SelectField', () => {
 
     expect(wrapper.find(ReactSelect).exists()).toBe(true)
 
-    expect(wrapper.findDataTest('select-field-icon-left').props()).toEqual({
+    expect(wrapper.findDataTestFirst('select-field-icon-left').props()).toEqual({
       'data-test': 'select-field-icon-left',
       containerClassName: styles.iconLeftContainer,
       icon: iconLeft,
diff --git a/packages/ui/__tests__/Toast.test.tsx b/packages/ui/__tests__/Toast.test.tsx
index dbfd7a3fa..803d9ac5c 100644
--- a/packages/ui/__tests__/Toast.test.tsx
+++ b/packages/ui/__tests__/Toast.test.tsx
@@ -47,7 +47,7 @@ describe('Toast', () => {
     expect(AlertElement.exists()).toBeTruthy()
     expect(AlertElement.findDataTest('toast-content').text()).toBe(content)
 
-    expect(wrapper.findDataTest('toast-icon').props()).toMatchObject({
+    expect(wrapper.findDataTestFirst('toast-icon').props()).toMatchObject({
       icon,
       ariaLabel,
     })
diff --git a/packages/ui/src/Button.tsx b/packages/ui/src/Button.tsx
index 17c59ca9e..37a94592b 100644
--- a/packages/ui/src/Button.tsx
+++ b/packages/ui/src/Button.tsx
@@ -10,6 +10,7 @@ import { LinkRel, LinkTarget } from './Link'
 import { Loader } from './Loader'
 
 import styles from './Button.module.scss'
+import { DataTestProp } from '@hazelcast/helpers'
 
 export type ButtonKind = 'primary' | 'secondary' | 'danger' | 'transparent'
 
@@ -105,7 +106,11 @@ export type ButtonTypeButtonProps = {
 
 export type ButtonTypeProps = ButtonTypeAnchorProps | ButtonTypeButtonProps
 
-export type ButtonProps<T = ButtonTypeProps> = ButtonCommonProps & T & ButtonAccessibleIconLeftProps & ButtonAccessibleIconRightProps
+export type ButtonProps<T = ButtonTypeProps> = ButtonCommonProps &
+  T &
+  ButtonAccessibleIconLeftProps &
+  ButtonAccessibleIconRightProps &
+  DataTestProp
 
 const capitalizeFirstCharacter = (str: string) => `${str[0].toUpperCase()}${str.slice(1)}`
 
diff --git a/packages/ui/src/Table/types.ts b/packages/ui/src/Table/types.ts
index 08b17a1e0..8bfe6a459 100644
--- a/packages/ui/src/Table/types.ts
+++ b/packages/ui/src/Table/types.ts
@@ -9,6 +9,8 @@ export type CellProps<D extends object> = {
   row: RowType<D>
 }
 
+export type ColumnId<D extends object> = Extract<keyof D, string>
+
 export type Column = {
   id: string
 }
@@ -17,7 +19,7 @@ export type ColumnType<D extends object> = {
   minWidth?: number
   maxWidth?: number
   Footer?: string | ((props: HeaderProps<D>) => JSX.Element | string)
-  accessor?: Extract<keyof D, string> | ((data: D, index: number) => string | number | JSX.Element)
+  accessor?: ColumnId<D> | ((data: D, index: number) => string | number | JSX.Element)
   Cell?: string | ((props: CellProps<D>) => JSX.Element)
   canHide?: boolean // true by default
   align?: 'left' | 'center' | 'right'
@@ -48,7 +50,7 @@ export type CellType<D extends object> = {
 }
 
 export type SortingRule<D extends object> = {
-  id: Extract<keyof D, string>
+  id: ColumnId<D>
   desc?: boolean | undefined
 }
 
@@ -56,14 +58,14 @@ export type InitialState = {
   paginationOptions?: { pageIndex?: number; pageSize?: number }
 }
 
-export type TableState = {
-  hiddenColumns: string[]
-  columnOrder: string[]
+export type TableState<D extends object> = {
+  hiddenColumns: ColumnId<D>[]
+  columnOrder: ColumnId<D>[]
   pageSize: number
   columnResizing: {
-    columnWidths: Record<string, number>
+    columnWidths: Record<ColumnId<D>, number>
   }
-  sortBy: { id: string; desc: boolean }[]
+  sortBy: SortingRule<D>[]
 }
 
 declare module '@tanstack/react-table' {
diff --git a/packages/ui/src/Table/useTrackTableState.ts b/packages/ui/src/Table/useTrackTableState.ts
index 43e5bbce1..e113c2994 100644
--- a/packages/ui/src/Table/useTrackTableState.ts
+++ b/packages/ui/src/Table/useTrackTableState.ts
@@ -2,9 +2,9 @@ import { useEffect, useRef } from 'react'
 import { Table } from '@tanstack/react-table'
 
 import { useRefValue } from '../hooks'
-import { TableState } from './types'
+import { ColumnId, SortingRule, TableState } from './types'
 
-export const useTrackTableState = <D>(tableInstance: Table<D>, onChange?: (state: TableState) => void) => {
+export const useTrackTableState = <D extends object>(tableInstance: Table<D>, onChange?: (state: TableState<D>) => void) => {
   const {
     columnOrder,
     columnVisibility,
@@ -19,15 +19,15 @@ export const useTrackTableState = <D>(tableInstance: Table<D>, onChange?: (state
     const cb = getOnChange()
     if (cb && skippedInitialStateChangeRef.current) {
       cb({
-        sortBy: sorting,
+        sortBy: sorting as SortingRule<D>[],
         pageSize,
-        columnOrder,
+        columnOrder: columnOrder as ColumnId<D>[],
         columnResizing: {
-          columnWidths: columnSizing,
+          columnWidths: columnSizing as Record<ColumnId<D>, number>,
         },
-        hiddenColumns: Object.entries(columnVisibility).reduce<string[]>((res, [id, visible]) => {
+        hiddenColumns: Object.entries(columnVisibility).reduce<ColumnId<D>[]>((res, [id, visible]) => {
           if (!visible) {
-            return [...res, id]
+            return [...res, id as ColumnId<D>]
           }
 
           return res