From cf08d5a56553969acda30abfd4f3b4bcab36243f Mon Sep 17 00:00:00 2001
From: Andre Wiggins <andrewiggins@live.com>
Date: Fri, 18 Nov 2022 16:37:24 -0800
Subject: [PATCH] Invoke setState callbacks setup in componentWillMount

Fixes #556
---
 src/diff/index.js                             |  3 +-
 .../lifecycles/componentWillMount.test.js     | 40 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/diff/index.js b/src/diff/index.js
index 599f26113b..242bf84581 100644
--- a/src/diff/index.js
+++ b/src/diff/index.js
@@ -110,6 +110,7 @@ export function diff(
 
 			oldProps = c.props;
 			oldState = c.state;
+			c._vnode = newVNode;
 
 			// Invoke pre-render lifecycle methods
 			if (isNew) {
@@ -146,7 +147,6 @@ export function diff(
 					c.state = c._nextState;
 					// More info about this here: https://gist.github.com/JoviDeCroock/bec5f2ce93544d2e6070ef8e0036e4e8
 					if (newVNode._original !== oldVNode._original) c._dirty = false;
-					c._vnode = newVNode;
 					newVNode._dom = oldVNode._dom;
 					newVNode._children = oldVNode._children;
 					newVNode._children.forEach(vnode => {
@@ -178,7 +178,6 @@ export function diff(
 
 			c.context = componentContext;
 			c.props = newProps;
-			c._vnode = newVNode;
 			c._parentDom = parentDom;
 
 			let renderHook = options._render,
diff --git a/test/browser/lifecycles/componentWillMount.test.js b/test/browser/lifecycles/componentWillMount.test.js
index 880803e84b..27b5c91212 100644
--- a/test/browser/lifecycles/componentWillMount.test.js
+++ b/test/browser/lifecycles/componentWillMount.test.js
@@ -1,4 +1,5 @@
 import { createElement, render, Component } from 'preact';
+import { setupRerender } from 'preact/test-utils';
 import { setupScratch, teardown } from '../../_util/helpers';
 
 /** @jsx createElement */
@@ -7,8 +8,12 @@ describe('Lifecycle methods', () => {
 	/** @type {HTMLDivElement} */
 	let scratch;
 
+	/** @type {() => void} */
+	let rerender;
+
 	beforeEach(() => {
 		scratch = setupScratch();
+		rerender = setupRerender();
 	});
 
 	afterEach(() => {
@@ -39,5 +44,40 @@ describe('Lifecycle methods', () => {
 
 			expect(componentState).to.deep.equal({ value: 1 });
 		});
+
+		it('should invoke setState callbacks when setState is called in componentWillMount', () => {
+			let componentState;
+			let callback = sinon.spy();
+
+			class Foo extends Component {
+				constructor(props) {
+					super(props);
+					this.state = {
+						value: 0
+					};
+				}
+				componentWillMount() {
+					this.setState({ value: 1 }, callback);
+					this.setState({ value: 2 }, () => {
+						callback();
+						this.setState({ value: 3 }, callback);
+					});
+				}
+				render() {
+					componentState = this.state;
+					return <div />;
+				}
+			}
+
+			render(<Foo />, scratch);
+
+			expect(componentState).to.deep.equal({ value: 2 });
+			expect(callback).to.have.been.calledTwice;
+
+			rerender();
+
+			expect(componentState).to.deep.equal({ value: 3 });
+			expect(callback).to.have.been.calledThrice;
+		});
 	});
 });