From 585e35b6b254aa6c4fb955d78985e6d7dba7c1fd Mon Sep 17 00:00:00 2001
From: Seyyed Morteza Moosavi <se.mo.moosavi@gmail.com>
Date: Wed, 18 Nov 2020 18:35:07 +0330
Subject: [PATCH] feat(fns-wire): support rewiring

---
 src/fn-wire/use-fn.test.ts  | 78 +++++++++++++++++++++++++++++++++++++
 src/fn-wire/use-fn.ts       |  2 -
 src/fn-wire/use-fns-wire.ts | 13 ++++++-
 3 files changed, 89 insertions(+), 4 deletions(-)

diff --git a/src/fn-wire/use-fn.test.ts b/src/fn-wire/use-fn.test.ts
index c4e1aaf..4fe9ca7 100644
--- a/src/fn-wire/use-fn.test.ts
+++ b/src/fn-wire/use-fn.test.ts
@@ -83,4 +83,82 @@ describe('useFn', () => {
     });
     expect(fn).toBeCalledWith(3);
   });
+  describe('when wire changed', () => {
+    it('should subscribe to the new wire', () => {
+      const fn = jest.fn();
+      const { result, rerender } = renderHook(
+        ({ firstWire }: { firstWire: boolean }) => {
+          const wire1 = useFnsWire<{ test(n: number): void }>(null);
+          const wire2 = useFnsWire<{ test(n: number): void }>(null);
+          useFn(firstWire ? wire1 : wire2, 'test', fn);
+          return { wire1, wire2 };
+        },
+        { initialProps: { firstWire: true } },
+      );
+      rerender({ firstWire: false });
+
+      act(() => {
+        result.current.wire2.fns.test(3);
+      });
+      expect(fn).toBeCalledWith(3);
+    });
+    it('should unsubscribe from first wire ', () => {
+      const fn = jest.fn();
+      const { result, rerender } = renderHook(
+        ({ firstWire }: { firstWire: boolean }) => {
+          const wire1 = useFnsWire<{ test(n: number): void }>(null);
+          const wire2 = useFnsWire<{ test(n: number): void }>(null);
+          useFn(firstWire ? wire1 : wire2, 'test', fn);
+          return { wire1, wire2 };
+        },
+        { initialProps: { firstWire: true } },
+      );
+      rerender({ firstWire: false });
+
+      act(() => {
+        result.current.wire1.fns.test(3);
+      });
+      expect(fn).not.toBeCalled();
+    });
+  });
+  describe('when up-link wire changed', () => {
+    it('should subscribe to the new wire', () => {
+      const fn = jest.fn();
+      const { result, rerender } = renderHook(
+        ({ firstWire }: { firstWire: boolean }) => {
+          const wire1 = useFnsWire<{ test(n: number): void }>(null);
+          const wire2 = useFnsWire<{ test(n: number): void }>(null);
+          const wire = useFnsWire(firstWire ? wire1 : wire2);
+          useFn(wire, 'test', fn);
+          return { wire1, wire2 };
+        },
+        { initialProps: { firstWire: true } },
+      );
+      rerender({ firstWire: false });
+
+      act(() => {
+        result.current.wire2.fns.test(3);
+      });
+      expect(fn).toBeCalledWith(3);
+    });
+    it('should unsubscribe from first wire ', () => {
+      const fn = jest.fn();
+      const { result, rerender } = renderHook(
+        ({ firstWire }: { firstWire: boolean }) => {
+          const wire1 = useFnsWire<{ test(n: number): void }>(null);
+          const wire2 = useFnsWire<{ test(n: number): void }>(null);
+          const wire = useFnsWire(firstWire ? wire1 : wire2);
+          useFn(wire, 'test', fn);
+          return { wire1, wire2 };
+        },
+        { initialProps: { firstWire: true } },
+      );
+      rerender({ firstWire: false });
+
+      act(() => {
+        result.current.wire1.fns.test(3);
+      });
+      expect(fn).not.toBeCalled();
+    });
+  });
 });
diff --git a/src/fn-wire/use-fn.ts b/src/fn-wire/use-fn.ts
index b2eaca9..dbc621b 100644
--- a/src/fn-wire/use-fn.ts
+++ b/src/fn-wire/use-fn.ts
@@ -1,6 +1,5 @@
 import { useEffect } from 'react';
 import { KeyOfMethods } from '../utils/type-utils';
-import { useStabilityGuard } from '../utils/use-stability-guard';
 import { FnsWire } from './fns-wire';
 
 /**
@@ -31,7 +30,6 @@ export function useFn<Fns, K extends KeyOfMethods<Fns>>(
   name: K,
   fn: Fns[K],
 ) {
-  useStabilityGuard(wire);
   useEffect(() => {
     return wire?.fn(name, fn);
   }, [wire, name, fn]);
diff --git a/src/fn-wire/use-fns-wire.ts b/src/fn-wire/use-fns-wire.ts
index 522a70b..7e5725b 100644
--- a/src/fn-wire/use-fns-wire.ts
+++ b/src/fn-wire/use-fns-wire.ts
@@ -1,11 +1,20 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { createFnsWire } from './create-fns-wire';
 import { FnsWire } from './fns-wire';
 
+const create = <Fns>(upLink: FnsWire<Fns> | null | undefined) => {
+  return createFnsWire(upLink || {});
+};
+
 export function useFnsWire<Fns = {}>(
   upLink: FnsWire<Fns> | null | undefined,
 ): FnsWire<Fns> {
-  const [[wire, connect]] = useState(() => createFnsWire(upLink || {}));
+  const [[wire, connect], set] = useState(() => create(upLink));
+  const lastUpLinkRef = useRef(upLink);
+  if (lastUpLinkRef.current !== upLink) {
+    lastUpLinkRef.current = upLink;
+    set(create(upLink));
+  }
   useEffect(() => {
     return connect();
   }, [connect]);