From 986fe7b8eb0dc5c59bc2a0525cdd7b83aeff982f Mon Sep 17 00:00:00 2001
From: zombiej <smith3816@gmail.com>
Date: Wed, 15 Sep 2021 14:39:01 +0800
Subject: [PATCH] fix: shouldCellUpdate should not ignore className

---
 src/Cell/index.tsx  |  4 +++-
 tests/Cell.spec.tsx | 28 ++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/Cell/index.tsx b/src/Cell/index.tsx
index dab5296b2..c6e3dc6d2 100644
--- a/src/Cell/index.tsx
+++ b/src/Cell/index.tsx
@@ -216,11 +216,13 @@ function Cell<RecordType extends DefaultRecordType>(
 const RefCell = React.forwardRef<any, CellProps<any>>(Cell);
 RefCell.displayName = 'Cell';
 
+const comparePropList: (keyof CellProps<any>)[] = ['expanded', 'className'];
+
 const MemoCell = React.memo(RefCell, (prev: CellProps<any>, next: CellProps<any>) => {
   if (next.shouldCellUpdate) {
     return (
       // Additional handle of expanded logic
-      prev.expanded === next.expanded &&
+      comparePropList.every(propName => prev[propName] === next[propName]) &&
       // User control update logic
       !next.shouldCellUpdate(next.record, prev.record)
     );
diff --git a/tests/Cell.spec.tsx b/tests/Cell.spec.tsx
index e37b9334a..84628a955 100644
--- a/tests/Cell.spec.tsx
+++ b/tests/Cell.spec.tsx
@@ -42,4 +42,32 @@ describe('Table.Cell', () => {
       expect(reRenderTime).toEqual(0);
     }
   });
+
+  it('shouldCellUpdate not block className', () => {
+    let reRenderTime = 0;
+
+    const getColumns = (props?: object) => [
+      {
+        shouldCellUpdate: (record, prevRecord) => prevRecord.key !== record.key,
+        dataIndex: 'key',
+        render: value => {
+          reRenderTime += 1;
+          return value;
+        },
+        ...props,
+      },
+    ];
+
+    const wrapper = mount(<Table data={[{ key: 'light' }]} columns={getColumns()} />);
+
+    // Update className should re-render
+    reRenderTime = 0;
+    for (let i = 0; i < 10; i += 1) {
+      wrapper.setProps({
+        columns: getColumns({ className: 'test' }),
+      });
+    }
+
+    expect(reRenderTime).toEqual(1);
+  });
 });