fix: Affix dynamic change target should work (#16146)

* fix target not refreshed

* add comment

components/affix/__tests__/Affix.test.js

@@ -1,6 +1,7 @@
 import React from 'react';
 import { mount } from 'enzyme';
 import Affix from '..';
+import { getObserverEntities } from '../utils';
 import Button from '../../button';
 
 const events = {};
@@ -128,14 +129,37 @@ describe('Affix Render', () => {
     expect(wrapper.instance().affix.state.affixStyle.top).toBe(10);
   });
 
-  it('updatePosition when target changed', () => {
-    const container = '<div id="mounter" />';
-    const getTarget = () => container;
-    wrapper = mount(<Affix target={getTarget} />);
-    wrapper.setProps({ target: null });
-    expect(wrapper.instance().state.status).toBe(0);
-    expect(wrapper.instance().state.affixStyle).toBe(undefined);
-    expect(wrapper.instance().state.placeholderStyle).toBe(undefined);
+  describe('updatePosition when target changed', () => {
+    it('function change', () => {
+      document.body.innerHTML = '<div id="mounter" />';
+      const container = document.querySelector('#id');
+      const getTarget = () => container;
+      wrapper = mount(<Affix target={getTarget} />);
+      wrapper.setProps({ target: null });
+      expect(wrapper.instance().state.status).toBe(0);
+      expect(wrapper.instance().state.affixStyle).toBe(undefined);
+      expect(wrapper.instance().state.placeholderStyle).toBe(undefined);
+    });
+
+    it('instance change', () => {
+      const getObserverLength = () => Object.keys(getObserverEntities()).length;
+
+      const container = document.createElement('div');
+      document.body.appendChild(container);
+      let target = container;
+
+      const originLength = getObserverLength();
+      const getTarget = () => target;
+      wrapper = mount(<Affix target={getTarget} />);
+      jest.runAllTimers();
+
+      expect(getObserverLength()).toBe(originLength + 1);
+      target = null;
+      wrapper.setProps({});
+      wrapper.update();
+      jest.runAllTimers();
+      expect(getObserverLength()).toBe(originLength);
+    });
   });
 
   it('updatePosition when size changed', () => {

components/affix/index.tsx

@@ -41,6 +41,8 @@ export interface AffixState {
   placeholderStyle?: React.CSSProperties;
   status: AffixStatus;
   lastAffix: boolean;
+
+  prevTarget: Window | HTMLElement | null;
 }
 
 class Affix extends React.Component<AffixProps, AffixState> {
@@ -51,6 +53,7 @@ class Affix extends React.Component<AffixProps, AffixState> {
   state: AffixState = {
     status: AffixStatus.None,
     lastAffix: false,
+    prevTarget: null,
   };
 
   placeholderNode: HTMLDivElement;
@@ -72,14 +75,22 @@ class Affix extends React.Component<AffixProps, AffixState> {
   }
 
   componentDidUpdate(prevProps: AffixProps) {
+    const { prevTarget } = this.state;
     const { target } = this.props;
-    if (prevProps.target !== target) {
+    let newTarget = null;
+    if (target) {
+      newTarget = target() || null;
+    }
+
+    if (prevTarget !== newTarget) {
       removeObserveTarget(this);
-      if (target) {
-        addObserveTarget(target(), this);
+      if (newTarget) {
+        addObserveTarget(newTarget, this);
         // Mock Event object.
         this.updatePosition({} as Event);
       }
+
+      this.setState({ prevTarget: newTarget });
     }
 
     if (

components/affix/utils.ts

@@ -20,6 +20,11 @@ interface ObserverEntity {
 
 let observerEntities: ObserverEntity[] = [];
 
+export function getObserverEntities() {
+  // Only used in test env. Can be removed if refactor.
+  return observerEntities;
+}
+
 export function addObserveTarget(target: HTMLElement | Window | null, affix: Affix): void {
   if (!target) return;