diff --git a/hooks/src/index.js b/hooks/src/index.js index d6c434e400..1f40937fbd 100644 --- a/hooks/src/index.js +++ b/hooks/src/index.js @@ -320,8 +320,11 @@ export function useImperativeHandle(ref, createHandle, args) { useLayoutEffect( () => { if (typeof ref == 'function') { - ref(createHandle()); - return () => ref(null); + const result = ref(createHandle()); + return () => { + ref(null); + if (result && typeof result == 'function') result(); + }; } else if (ref) { ref.current = createHandle(); return () => (ref.current = null); diff --git a/hooks/test/browser/useImperativeHandle.test.js b/hooks/test/browser/useImperativeHandle.test.js index ac71065696..f05df6212c 100644 --- a/hooks/test/browser/useImperativeHandle.test.js +++ b/hooks/test/browser/useImperativeHandle.test.js @@ -35,6 +35,30 @@ describe('useImperativeHandle', () => { expect(ref.current.test()).to.equal('test'); }); + it('Calls ref unmounting function', () => { + let ref; + const unmount = sinon.spy(); + + function Comp() { + useImperativeHandle( + r => { + ref = r; + return unmount; + }, + () => ({ test: () => 'test' }), + [] + ); + return

Test

; + } + + render(, scratch); + expect(ref).to.have.property('test'); + expect(ref.test()).to.equal('test'); + render(null, scratch); + expect(unmount).to.be.calledOnce; + expect(ref).to.equal(null); + }); + it('calls createHandle after every render by default', () => { let ref, createHandleSpy = sinon.spy();