Support ref unmounting on imperative handles (#4625)

Author: JoviDeCroock

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);

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 <p>Test</p>;
+		}
+
+		render(<Comp />, 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();