add new tests to partitionByKey

Author: voliva

packages/utils/src/partitionByKey.test.ts

@@ -1,5 +1,5 @@
-import { NEVER, Observable, Subject } from "rxjs"
-import { switchMap, take } from "rxjs/operators"
+import { concat, NEVER, Observable, of, Subject } from "rxjs"
+import { catchError, switchMap, take } from "rxjs/operators"
 import { TestScheduler } from "rxjs/testing"
 import { partitionByKey } from "./"
 
@@ -58,6 +58,33 @@ describe("partitionByKey", () => {
         })
       })
     })
+
+    it("keeps the existing keys alive when the source completes", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const source = cold("-ab-|")
+        const a = cold("      --1---2-|")
+        const b = cold("       ---|")
+        const expectedStr = "efg--h---(i|)"
+        const innerStreams = { a, b }
+        const [, result] = partitionByKey(
+          source,
+          (v) => v,
+          (v$) =>
+            v$.pipe(
+              take(1),
+              switchMap((v) => innerStreams[v]),
+            ),
+        )
+
+        expectObservable(result).toBe(expectedStr, {
+          e: [],
+          f: ["a"],
+          g: ["a", "b"],
+          h: ["a"],
+          i: [],
+        })
+      })
+    })
   })
 
   describe("getInstance$", () => {
@@ -112,5 +139,95 @@ describe("partitionByKey", () => {
       expect(lateNext).toHaveBeenCalledTimes(1)
       expect(lateNext).toHaveBeenCalledWith(1)
     })
+
+    it("unsubscribes from all streams when refcount reaches 0", () => {
+      let innerSubs = 0
+      const inner = new Observable<number>(() => {
+        innerSubs++
+        return () => {
+          innerSubs--
+        }
+      })
+
+      const sourceSubject = new Subject<number>()
+      let sourceSubs = 0
+      const source = new Observable<number>((obs) => {
+        sourceSubs++
+        sourceSubject.subscribe(obs)
+        return () => {
+          sourceSubs--
+        }
+      })
+
+      const [getObs] = partitionByKey(
+        source,
+        (v) => v,
+        () => inner,
+      )
+      const observable = getObs(1)
+
+      expect(sourceSubs).toBe(0)
+      expect(innerSubs).toBe(0)
+
+      const sub1 = observable.subscribe()
+
+      expect(sourceSubs).toBe(1)
+      expect(innerSubs).toBe(0)
+
+      sourceSubject.next(1)
+
+      expect(sourceSubs).toBe(1)
+      expect(innerSubs).toBe(1)
+
+      const sub2 = observable.subscribe()
+
+      expect(sourceSubs).toBe(1)
+      expect(innerSubs).toBe(1)
+
+      sub1.unsubscribe()
+
+      expect(sourceSubs).toBe(1)
+      expect(innerSubs).toBe(1)
+
+      sub2.unsubscribe()
+
+      expect(sourceSubs).toBe(0)
+      expect(innerSubs).toBe(0)
+    })
+
+    it("emits a complete on the inner observable when the source completes", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const source = cold("-ab-a-|")
+        const expectA = "    -a--a-(c|)"
+        const expectB = "    --b---(c|)"
+
+        const [getInstance$] = partitionByKey(
+          source,
+          (v) => v,
+          (v$) => concat(v$, ["c"]),
+        )
+
+        expectObservable(getInstance$("a")).toBe(expectA)
+        expectObservable(getInstance$("b")).toBe(expectB)
+      })
+    })
+
+    // Do we want this behaviour?
+    it("emits an error when the source errors", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const source = cold("-ab-a-#")
+        const expectA = "    -a--a-#"
+        const expectB = "    --b---#"
+
+        const [getInstance$] = partitionByKey(
+          source,
+          (v) => v,
+          (v$) => v$.pipe(catchError(() => of("e"))),
+        )
+
+        expectObservable(getInstance$("a")).toBe(expectA)
+        expectObservable(getInstance$("b")).toBe(expectB)
+      })
+    })
   })
 })