From c7c886fda6f46cbb4d550e92e3f1294d2af8b8f7 Mon Sep 17 00:00:00 2001
From: afc163 <afc163@gmail.com>
Date: Sat, 3 Dec 2016 16:19:22 +0800
Subject: [PATCH] Improve Affix target demo code style, close #4037

---
 components/affix/demo/target.md | 30 +++++++++++++-----------------
 components/affix/index.tsx      | 10 +++++++++-
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/components/affix/demo/target.md b/components/affix/demo/target.md
index 52bbc6d71d..d78ab9100e 100644
--- a/components/affix/demo/target.md
+++ b/components/affix/demo/target.md
@@ -16,36 +16,32 @@ Set a `target` for 'Affix', which is listen to scroll event of target element (d
 ````jsx
 import { Affix, Button } from 'antd';
 
-const Demo = () => {
-  return (
-    <div className="view-port">
-      <div id="scrollable-container">
+class Demo extends React.Component {
+  render() {
+    return (
+      <div className="scrollable-container" ref={(node) => { this.container = node; }}>
         <div className="background">
-          <br />
-          <br />
-          <br />
-          <Affix target={() => document.getElementById('scrollable-container')} offsetTop={20}>
-            <Button type="primary">Fixed at the top of container</Button>
+          <Affix target={() => this.container}>
+            <Button type="primary">
+              Fixed at the top of container
+            </Button>
           </Affix>
         </div>
       </div>
-    </div>
-  );
-};
+    );
+  }
+}
 
 ReactDOM.render(<Demo />, mountNode);
 ````
 
 <style>
-#components-affix-demo-target .view-port {
+#components-affix-demo-target .scrollable-container {
   height: 100px;
-  overflow: hidden;
-}
-#components-affix-demo-target #scrollable-container {
-  height: 100%;
   overflow-y: scroll;
 }
 #components-affix-demo-target .background {
+  padding-top: 60px;
   height: 300px;
   background-image: url('https://zos.alipayobjects.com/rmsportal/RmjwQiJorKyobvI.jpg');
 }
diff --git a/components/affix/index.tsx b/components/affix/index.tsx
index 6642c71eff..15a0e728bb 100644
--- a/components/affix/index.tsx
+++ b/components/affix/index.tsx
@@ -64,6 +64,7 @@ export default class Affix extends React.Component<AffixProps, any> {
 
   scrollEvent: any;
   resizeEvent: any;
+  timeout: any;
   refs: {
     fixedNode: HTMLElement;
   };
@@ -172,7 +173,10 @@ export default class Affix extends React.Component<AffixProps, any> {
 
   componentDidMount() {
     const target = this.props.target || getDefaultTarget;
-    this.setTargetEventListeners(target);
+    // Wait for parent component ref has its value
+    this.timeout = setTimeout(() => {
+      this.setTargetEventListeners(target);
+    });
   }
 
   componentWillReceiveProps(nextProps) {
@@ -187,10 +191,14 @@ export default class Affix extends React.Component<AffixProps, any> {
 
   componentWillUnmount() {
     this.clearScrollEventListeners();
+    clearTimeout(this.timeout);
   }
 
   setTargetEventListeners(getTarget) {
     const target = getTarget();
+    if (!target) {
+      return;
+    }
     this.scrollEvent = addEventListener(target, 'scroll', this.updatePosition);
     this.resizeEvent = addEventListener(target, 'resize', this.updatePosition);
   }