Add test for unwritable streams in demux

tests/compose.spec.ts

@@ -219,10 +219,10 @@ test("compose() should emit drain event ~rate * highWaterMark ms for every write
 
         const input = [
             { key: "a", mapped: [] },
-            { key: "b", mapped: [] },
-            { key: "c", mapped: [] },
-            { key: "d", mapped: [] },
-            { key: "e", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
         ];
 
         let start = performance.now();

tests/demux.spec.ts

@@ -17,6 +17,8 @@ test.cb("demux() constructor should be called once per key", t => {
         { key: "b", visited: [] },
         { key: "a", visited: [] },
         { key: "c", visited: [] },
+        { key: "a", visited: [] },
+        { key: "b", visited: [] },
     ];
     const construct = sinon.spy((destKey: string) => {
         const dest = map((chunk: Test) => {
@@ -48,6 +50,8 @@ test.cb("demux() constructor should be called once per key using keyBy", t => {
         { key: "b", visited: [] },
         { key: "a", visited: [] },
         { key: "c", visited: [] },
+        { key: "a", visited: [] },
+        { key: "b", visited: [] },
     ];
 
     const construct = sinon.spy((destKey: string) => {
@@ -134,15 +138,15 @@ test.cb("should emit errors", t => {
     demuxed.end();
 });
 
-test("compose() should emit drain event ~rate * highWaterMark ms for every write that causes backpressure", async t => {
+test("demux() should emit drain event ~rate * highWaterMark ms for every write that causes backpressure", t => {
     t.plan(7);
+    interface Chunk {
+        key: string;
+        mapped: number[];
+    }
     const highWaterMark = 5;
     const _rate = 25;
     return new Promise(async (resolve, reject) => {
-        interface Chunk {
-            key: string;
-            mapped: number[];
-        }
         const sink = new Writable({
             objectMode: true,
             write(chunk, encoding, cb) {
@@ -222,7 +226,8 @@ test("demux() should emit one drain event when writing 6 items with highWaterMar
             write(chunk, encoding, cb) {
                 cb();
                 t.pass();
-                if (chunk.key === "f") {
+                pendingReads--;
+                if (pendingReads === 0) {
                     resolve();
                 }
             },
@@ -260,12 +265,13 @@ test("demux() should emit one drain event when writing 6 items with highWaterMar
 
         const input = [
             { key: "a", mapped: [] },
-            { key: "b", mapped: [] },
-            { key: "c", mapped: [] },
-            { key: "d", mapped: [] },
-            { key: "e", mapped: [] },
-            { key: "f", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
         ];
+        let pendingReads = input.length;
 
         for (const item of input) {
             const res = _demux.write(item);
@@ -354,6 +360,7 @@ test.cb(
         const start = performance.now();
     },
 );
+
 test.cb(
     "demux() should emit drain event immediately when second stream is bottleneck",
     t => {
@@ -431,6 +438,79 @@ test.cb(
     },
 );
 
+test("demux() should only emit drain event when all streams are writable", t => {
+    t.plan(1);
+    const highWaterMark = 2;
+    interface Chunk {
+        key: string;
+        mapped: number[];
+    }
+    return new Promise(async (resolve, reject) => {
+        const sink = new Writable({
+            objectMode: true,
+            write(chunk, encoding, cb) {
+                cb();
+                pendingReads--;
+                if (chunk.key === "b") {
+                    expect(performance.now() - start).to.be.greaterThan(150);
+                    t.pass();
+                }
+                if (pendingReads === 0) {
+                    resolve();
+                }
+            },
+        });
+        const construct = (destKey: string) => {
+            const first = map(
+                (chunk: Chunk) => {
+                    chunk.mapped.push(1);
+                    return chunk;
+                },
+                { objectMode: true },
+            );
+
+            const second = map(
+                async (chunk: Chunk) => {
+                    await sleep(50);
+                    chunk.mapped.push(2);
+                    return chunk;
+                },
+                { objectMode: true, highWaterMark: 1 },
+            );
+
+            first.pipe(second).pipe(sink);
+            return first;
+        };
+        const _demux = demux(
+            construct,
+            { key: "key" },
+            {
+                objectMode: true,
+                highWaterMark,
+            },
+        );
+        _demux.on("error", err => {
+            reject();
+        });
+
+        const input = [
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "c", mapped: [] },
+            { key: "c", mapped: [] },
+            { key: "b", mapped: [] }, // should only be recieved after a becomes writable
+        ];
+
+        let pendingReads = input.length;
+        let start = performance.now();
+        for (const item of input) {
+            const res = _demux.write(item);
+            if (!res) {
+                await sleep(50);
+            }
+        }
+    });
+});
 test("demux() should emit drain event and first should contain up to highWaterMark items in readable state when second is bottleneck", t => {
     t.plan(6);
     const highWaterMark = 5;