From 4b0dba5db8659512dc3ab4229e66c6aee7d1887a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Oliva?= <olivarra1@gmail.com>
Date: Mon, 9 Nov 2020 22:13:50 +0100
Subject: [PATCH] fix(core/subscribe): stop rendering children when source$
 switches

---
 packages/core/src/Subscribe.test.tsx | 85 +++++++++++++++++++++++++++-
 packages/core/src/Subscribe.tsx      | 12 ++--
 2 files changed, 90 insertions(+), 7 deletions(-)

diff --git a/packages/core/src/Subscribe.test.tsx b/packages/core/src/Subscribe.test.tsx
index 6e68efb..7d29f19 100644
--- a/packages/core/src/Subscribe.test.tsx
+++ b/packages/core/src/Subscribe.test.tsx
@@ -1,7 +1,7 @@
-import React from "react"
 import { render } from "@testing-library/react"
-import { Observable } from "rxjs"
-import { Subscribe, bind } from "./"
+import React, { useLayoutEffect, useState } from "react"
+import { defer, Observable, of } from "rxjs"
+import { bind, Subscribe } from "./"
 
 describe("Subscribe", () => {
   it("subscribes to the provided observable and remains subscribed until it's unmounted", () => {
@@ -31,4 +31,83 @@ describe("Subscribe", () => {
     unmount()
     expect(nSubscriptions).toBe(0)
   })
+
+  it("doesn't render its content until it has subscribed to a the source", () => {
+    let nSubscriptions = 0
+    let errored = false
+    const [useInstance, instance$] = bind((id: number) => {
+      if (id === 0) {
+        return of(0)
+      }
+      return defer(() => {
+        nSubscriptions++
+        return of(1)
+      })
+    })
+
+    const Child = ({ id }: { id: number }) => {
+      const value = useInstance(id)
+
+      if (id !== 0 && nSubscriptions === 0) {
+        errored = true
+      }
+
+      return <>{value}</>
+    }
+    const { rerender } = render(
+      <Subscribe source$={instance$(0)}>
+        <Child id={0} />
+      </Subscribe>,
+    )
+    expect(nSubscriptions).toBe(0)
+    expect(errored).toBe(false)
+
+    rerender(
+      <Subscribe source$={instance$(1)}>
+        <Child id={1} />
+      </Subscribe>,
+    )
+    expect(nSubscriptions).toBe(1)
+    expect(errored).toBe(false)
+  })
+
+  it("prevents the issue of stale data when switching keys", () => {
+    const [useInstance, instance$] = bind((id: number) => of(id))
+
+    const Child = ({
+      id,
+      initialValue,
+    }: {
+      id: number
+      initialValue: number
+    }) => {
+      const [value] = useState(initialValue)
+
+      return (
+        <>
+          <div data-testid="id">{id}</div>
+          <div data-testid="value">{value}</div>
+        </>
+      )
+    }
+
+    const Parent = ({ id }: { id: number }) => {
+      const value = useInstance(id)
+
+      return <Child key={id} id={id} initialValue={value} />
+    }
+    const { rerender, getByTestId } = render(
+      <Subscribe source$={instance$(0)}>
+        <Parent id={0} />
+      </Subscribe>,
+    )
+
+    rerender(
+      <Subscribe source$={instance$(1)}>
+        <Parent id={1} />
+      </Subscribe>,
+    )
+    expect(getByTestId("id").textContent).toBe("1")
+    expect(getByTestId("value").textContent).toBe("1")
+  })
 })
diff --git a/packages/core/src/Subscribe.tsx b/packages/core/src/Subscribe.tsx
index e44c2a6..7f151dc 100644
--- a/packages/core/src/Subscribe.tsx
+++ b/packages/core/src/Subscribe.tsx
@@ -22,9 +22,9 @@ export const Subscribe: React.FC<{
   const [mounted, setMounted] = useState(() => {
     try {
       ;(source$ as any).gV()
-      return 1
+      return source$
     } catch (e) {
-      return e.then ? 1 : 0
+      return e.then ? source$ : null
     }
   })
   useLayoutEffect(() => {
@@ -33,11 +33,15 @@ export const Subscribe: React.FC<{
         throw e
       }),
     )
-    setMounted(1)
+    setMounted(source$)
     return () => {
       subscription.unsubscribe()
     }
   }, [source$])
   const fBack = fallback || null
-  return <Suspense fallback={fBack}>{mounted ? children : <Throw />}</Suspense>
+  return (
+    <Suspense fallback={fBack}>
+      {mounted === source$ ? children : <Throw />}
+    </Suspense>
+  )
 }