From 45f18a28852296f5635094a0eac254131c3ebb7c Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Sun, 18 Jul 2021 09:57:49 +0900
Subject: [PATCH] chore(utils): refactor atomWithObservable (#599)

---
 src/utils/atomWithObservable.ts         | 80 +++++++++++++------------
 tests/utils/atomWithObservable.test.tsx | 69 ++++++++++++++++++---
 2 files changed, 104 insertions(+), 45 deletions(-)

diff --git a/src/utils/atomWithObservable.ts b/src/utils/atomWithObservable.ts
index ae9b1bcd88..1a70c996d4 100644
--- a/src/utils/atomWithObservable.ts
+++ b/src/utils/atomWithObservable.ts
@@ -34,69 +34,73 @@ export function atomWithObservable<TData>(
 export function atomWithObservable<TData>(
   createObservable: (get: Getter) => ObservableLike<TData> | SubjectLike<TData>
 ) {
-  type Result = { data: TData } | { error: unknown }
   const observableResultAtom = atom((get) => {
-    let resolve: ((result: Result) => void) | null = null
+    let settlePromise: ((data: TData | null, err?: unknown) => void) | null =
+      null
     let observable = createObservable(get)
-    if (observable[Symbol.observable]) {
-      observable = (
-        observable[Symbol.observable] as () => ObservableLike<TData>
-      )()
+    const returnsItself = observable[Symbol.observable]
+    if (returnsItself) {
+      observable = returnsItself()
     }
-    const resultAtom = atom<Result | Promise<Result>>(
-      new Promise<Result>((r) => {
-        resolve = r
+    const dataAtom = atom<TData | Promise<TData>>(
+      new Promise<TData>((resolve, reject) => {
+        settlePromise = (data, err) => {
+          if (err) {
+            reject(err)
+          } else {
+            resolve(data as TData)
+          }
+        }
       })
     )
-    resultAtom.scope = observableAtom.scope
-    let setResult: (result: Result) => void = () => {
+    dataAtom.scope = observableAtom.scope
+    let setData: (data: TData | Promise<TData>) => void = () => {
       throw new Error('setting data without mount')
     }
     const dataListener = (data: TData) => {
-      if (resolve) {
-        resolve({ data })
-        resolve = null
+      if (settlePromise) {
+        settlePromise(data)
+        settlePromise = null
+        if (subscription && !setData) {
+          subscription.unsubscribe()
+          subscription = null
+        }
       } else {
-        setResult({ data })
+        setData(data)
       }
     }
     const errorListener = (error: unknown) => {
-      if (resolve) {
-        resolve({ error })
-        resolve = null
+      if (settlePromise) {
+        settlePromise(null, error)
+        settlePromise = null
+        if (subscription && !setData) {
+          subscription.unsubscribe()
+          subscription = null
+        }
       } else {
-        setResult({ error })
+        setData(Promise.reject<TData>(error))
       }
     }
     let subscription: Subscription | null = null
-    subscription = observable.subscribe((data) => {
-      dataListener(data)
-      if (subscription && !setResult) {
-        subscription.unsubscribe()
-        subscription = null
-      }
-    }, errorListener)
-    if (!resolve) {
+    subscription = observable.subscribe(dataListener, errorListener)
+    if (!settlePromise) {
       subscription.unsubscribe()
       subscription = null
     }
-    resultAtom.onMount = (update) => {
-      setResult = update
-      if (!subscription)
+    dataAtom.onMount = (update) => {
+      setData = update
+      if (!subscription) {
         subscription = observable.subscribe(dataListener, errorListener)
-      return () => subscription && subscription.unsubscribe()
+      }
+      return () => subscription?.unsubscribe()
     }
-    return { resultAtom, observable }
+    return { dataAtom, observable }
   })
   const observableAtom = atom(
     (get) => {
       observableResultAtom.scope = observableAtom.scope
-      const { resultAtom } = get(observableResultAtom)
-      const result = get(resultAtom)
-      if ('error' in result) {
-        throw result.error
-      }
-      return result.data
+      const { dataAtom } = get(observableResultAtom)
+      return get(dataAtom)
     },
     (get, _set, data: TData) => {
       observableResultAtom.scope = observableAtom.scope
diff --git a/tests/utils/atomWithObservable.test.tsx b/tests/utils/atomWithObservable.test.tsx
index 55dcd2aa07..ee0590dbf0 100644
--- a/tests/utils/atomWithObservable.test.tsx
+++ b/tests/utils/atomWithObservable.test.tsx
@@ -1,4 +1,4 @@
-import { FC, Suspense } from 'react'
+import { FC, Suspense, Component } from 'react'
 import { fireEvent, render } from '@testing-library/react'
 import { getTestProvider } from '../testUtils'
 import { useAtom } from '../../src/index'
@@ -7,6 +7,31 @@ import { Observable, Subject } from 'rxjs'
 
 const Provider = getTestProvider()
 
+class ErrorBoundary extends Component<
+  { message?: string },
+  { hasError: boolean }
+> {
+  constructor(props: { message?: string }) {
+    super(props)
+    this.state = { hasError: false }
+  }
+  static getDerivedStateFromError() {
+    return { hasError: true }
+  }
+  render() {
+    return this.state.hasError ? (
+      <div>
+        <div>{this.props.message || 'errored'}</div>
+        <button onClick={() => this.setState({ hasError: false })}>
+          retry
+        </button>
+      </div>
+    ) : (
+      this.props.children
+    )
+  }
+}
+
 it('count state', async () => {
   const observableAtom = atomWithObservable(
     () =>
@@ -16,14 +41,14 @@ it('count state', async () => {
   )
 
   const Counter: FC = () => {
-    const [state$] = useAtom(observableAtom)
+    const [state] = useAtom(observableAtom)
 
-    return <>count: {state$}</>
+    return <>count: {state}</>
   }
 
   const { findByText } = render(
     <Provider>
-      <Suspense fallback="loading...">
+      <Suspense fallback="loading">
         <Counter />
       </Suspense>
     </Provider>
@@ -46,11 +71,11 @@ it('writable count state', async () => {
   })
 
   const Counter: FC = () => {
-    const [state$, dispatch] = useAtom(observableAtom)
+    const [state, dispatch] = useAtom(observableAtom)
 
     return (
       <>
-        count: {state$}
+        count: {state}
         <button onClick={() => dispatch(9)}>button</button>
       </>
     )
@@ -58,14 +83,44 @@ it('writable count state', async () => {
 
   const { findByText, getByText } = render(
     <Provider>
-      <Suspense fallback="loading...">
+      <Suspense fallback="loading">
         <Counter />
       </Suspense>
     </Provider>
   )
 
+  await findByText('loading')
   await findByText('count: 1')
 
   fireEvent.click(getByText('button'))
   await findByText('count: 9')
 })
+
+// FIXME we would like to support retry
+it.skip('count state with error', async () => {
+  const myObservable = new Observable<number>((subscriber) => {
+    subscriber.error('err1')
+    subscriber.next(1)
+  })
+  const observableAtom = atomWithObservable(() => myObservable)
+
+  const Counter: React.FC = () => {
+    const [state] = useAtom(observableAtom)
+
+    return <div>count: {state}</div>
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <ErrorBoundary>
+        <Suspense fallback="loading">
+          <Counter />
+        </Suspense>
+      </ErrorBoundary>
+    </Provider>
+  )
+
+  await findByText('errored')
+  fireEvent.click(getByText('retry'))
+  await findByText('count: 1')
+})