From a3adfa6f2544e020c135df820e472f3363a2b072 Mon Sep 17 00:00:00 2001
From: oltrep <34498770+oltrep@users.noreply.github.com>
Date: Tue, 21 Sep 2021 12:19:11 -0700
Subject: [PATCH] Fix bug in context consumers willUnmount (#320)

---
 CHANGELOG.md               |  2 +
 src/createContext.lua      |  1 +
 src/createContext.spec.lua | 97 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5d7b70bf..844a4a2d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## Unreleased Changes
 
+* Fixed `Listeners can only be disconnected once` from context consumers. ([#320](https://github.com/Roblox/roact/pull/320))
+
 ## [1.4.1](https://github.com/Roblox/roact/releases/tag/v1.4.1) (August 12th, 2021)
 * Fixed a bug where the Roact tree could get into a broken state when using callbacks passed to a child component. Updated the tempFixUpdateChildrenReEntrancy config value to also handle this case. ([#315](https://github.com/Roblox/roact/pull/315))
 * Fixed forwardRef description ([#312](https://github.com/Roblox/roact/pull/312)).
diff --git a/src/createContext.lua b/src/createContext.lua
index b21635ef..e4dcae15 100644
--- a/src/createContext.lua
+++ b/src/createContext.lua
@@ -119,6 +119,7 @@ local function createConsumer(context)
 	function Consumer:willUnmount()
 		if self.disconnect ~= nil then
 			self.disconnect()
+			self.disconnect = nil
 		end
 	end
 
diff --git a/src/createContext.spec.lua b/src/createContext.spec.lua
index 432d39d4..4b436b86 100644
--- a/src/createContext.spec.lua
+++ b/src/createContext.spec.lua
@@ -1,4 +1,6 @@
 return function()
+	local ReplicatedStorage = game:GetService("ReplicatedStorage")
+
 	local Component = require(script.Parent.Component)
 	local NoopRenderer = require(script.Parent.NoopRenderer)
 	local Children = require(script.Parent.PropMarkers.Children)
@@ -10,6 +12,9 @@ return function()
 
 	local noopReconciler = createReconciler(NoopRenderer)
 
+	local RobloxRenderer = require(script.Parent.RobloxRenderer)
+	local robloxReconciler = createReconciler(RobloxRenderer)
+
 	it("should return a table", function()
 		local context = createContext("Test")
 		expect(context).to.be.ok()
@@ -301,4 +306,96 @@ return function()
 			expect(observedB).to.equal(true)
 		end)
 	end)
+
+	-- issue https://github.com/Roblox/roact/issues/319
+	it("does not throw if willUnmount is called twice on a context consumer", function()
+		local context = createContext({})
+
+		local LowestComponent = Component:extend("LowestComponent")
+		function LowestComponent:init()
+		end
+
+		function LowestComponent:render()
+			return createElement("Frame")
+		end
+
+		function LowestComponent:didMount()
+			self.props.onDidMountCallback()
+		end
+
+		local FirstComponent = Component:extend("FirstComponent")
+		function FirstComponent:init()
+		end
+
+		function FirstComponent:render()
+			return createElement(context.Consumer, {
+				render = function()
+					return createElement("TextLabel")
+				end,
+			})
+		end
+
+		local ChildComponent = Component:extend("ChildComponent")
+
+		function ChildComponent:init()
+			self:setState({ firstTime = true })
+		end
+
+		local childCallback
+
+		function ChildComponent:render()
+			if self.state.firstTime then
+				return createElement(FirstComponent)
+			end
+
+			return createElement(LowestComponent, {
+				onDidMountCallback = self.props.onDidMountCallback
+			})
+		end
+
+		function ChildComponent:didMount()
+			childCallback = function()
+				self:setState({ firstTime = false })
+			end
+		end
+
+		local ParentComponent = Component:extend("ParentComponent")
+
+		local didMountCallbackCalled = 0
+
+		function ParentComponent:init()
+			self:setState({ count = 1 })
+
+			self.onDidMountCallback = function()
+				didMountCallbackCalled = didMountCallbackCalled + 1
+				if self.state.count < 5 then
+					self:setState({ count = self.state.count + 1 })
+				end
+			end
+		end
+
+		function ParentComponent:render()
+			return createElement("Frame", {}, {
+				Provider = createElement(context.Provider, {
+					value = {},
+				}, {
+					ChildComponent = createElement(ChildComponent, {
+						count = self.state.count,
+						onDidMountCallback = self.onDidMountCallback,
+					}),
+				})
+			})
+		end
+
+		local parent = Instance.new("ScreenGui")
+		parent.Parent = ReplicatedStorage
+
+		local hostKey = "Some Key"
+		robloxReconciler.mountVirtualNode(createElement(ParentComponent), parent, hostKey)
+
+		expect(function()
+			-- calling setState on ChildComponent will trigger `willUnmount` multiple times
+			childCallback()
+		end).never.to.throw()
+	end)
 end
\ No newline at end of file