Add tests for demux

tests/compose.spec.ts

@@ -1,6 +1,6 @@
 const test = require("ava");
 const { expect } = require("chai");
-const { compose, composeDuplex, map, rate } = require("../src");
+const { compose, map } = require("../src");
 const { sleep } = require("../src/helpers");
 import { performance } from "perf_hooks";
 
@@ -308,12 +308,12 @@ test.cb(
     t => {
         t.plan(6);
         interface Chunk {
-            index: number;
-            mapped: string[];
+            key: string;
+            mapped: number[];
         }
         const first = map(
             async (chunk: Chunk) => {
-                chunk.mapped.push("first");
+                chunk.mapped.push(1);
                 return chunk;
             },
             {
@@ -324,9 +324,10 @@ test.cb(
         const second = map(
             async (chunk: Chunk) => {
                 pendingReads--;
-                await sleep(500);
+                await sleep(200);
+                expect(second._writableState.length).to.be.equal(1);
                 expect(first._readableState.length).to.equal(pendingReads);
-                chunk.mapped.push("second");
+                chunk.mapped.push(2);
                 return chunk;
             },
             { objectMode: true, highWaterMark: 1 },
@@ -342,27 +343,25 @@ test.cb(
 
         composed.on("drain", () => {
             expect(composed._writableState.length).to.be.equal(0);
-            expect(performance.now() - start).to.be.lessThan(100);
+            expect(performance.now() - start).to.be.lessThan(50);
             t.pass();
         });
 
         composed.on("data", (chunk: Chunk) => {
             // Since second is bottleneck, composed will write into first immediately. Buffer should be empty.
             expect(composed._writableState.length).to.be.equal(0);
-            expect(chunk.mapped.length).to.equal(2);
-            expect(chunk.mapped).to.deep.equal(["first", "second"]);
             t.pass();
-            if (chunk.index === 5) {
+            if (chunk.key === "e") {
                 t.end();
             }
         });
 
         const input = [
-            { index: 1, mapped: [] },
-            { index: 2, mapped: [] },
-            { index: 3, mapped: [] },
-            { index: 4, mapped: [] },
-            { index: 5, mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "b", mapped: [] },
+            { key: "c", mapped: [] },
+            { key: "d", mapped: [] },
+            { key: "e", mapped: [] },
         ];
         let pendingReads = input.length;
 

tests/demux.spec.ts

@@ -1,53 +1,38 @@
 import test from "ava";
 import { expect } from "chai";
-import { demux, map } from "../src";
+const { demux, map } = require("../src");
 import { Writable } from "stream";
+const sinon = require("sinon");
+const { sleep } = require("../src/helpers");
+import { performance } from "perf_hooks";
 
 interface Test {
     key: string;
-    val: number;
+    visited: number[];
 }
-test.cb("should spread per key", t => {
-    t.plan(5);
+test.cb("demux() constructor should be called once per key", t => {
+    t.plan(1);
     const input = [
-        { key: "a", val: 1 },
-        { key: "b", val: 2 },
-        { key: "a", val: 3 },
-        { key: "c", val: 4 },
+        { key: "a", visited: [] },
+        { key: "b", visited: [] },
+        { key: "a", visited: [] },
+        { key: "c", visited: [] },
     ];
-    const results = [
-        { key: "a", val: 2 },
-        { key: "b", val: 3 },
-        { key: "a", val: 4 },
-        { key: "c", val: 5 },
-    ];
-    const destinationStreamKeys = [];
-    let i = 0;
-    const sink = new Writable({
-        objectMode: true,
-        write(chunk, enc, cb) {
-            expect(results).to.deep.include(chunk);
-            expect(input).to.not.deep.include(chunk);
-            t.pass();
-            cb();
-        },
-    });
-    const construct = (destKey: string) => {
-        destinationStreamKeys.push(destKey);
+    const construct = sinon.spy((destKey: string) => {
         const dest = map((chunk: Test) => {
-            return {
-                ...chunk,
-                val: chunk.val + 1,
-            };
+            chunk.visited.push(1);
+            return chunk;
         });
 
-        dest.pipe(sink);
         return dest;
-    };
+    });
 
     const demuxed = demux(construct, { key: "key" }, { objectMode: true });
+
     demuxed.on("finish", () => {
-        expect(destinationStreamKeys).to.deep.equal(["a", "b", "c"]);
+        expect(construct.withArgs("a").callCount).to.equal(1);
+        expect(construct.withArgs("b").callCount).to.equal(1);
+        expect(construct.withArgs("c").callCount).to.equal(1);
         t.pass();
         t.end();
     });
@@ -56,50 +41,34 @@ test.cb("should spread per key", t => {
     demuxed.end();
 });
 
-test.cb("should spread per key using keyBy", t => {
-    t.plan(5);
+test.cb("demux() constructor should be called once per key using keyBy", t => {
+    t.plan(1);
     const input = [
-        { key: "a", val: 1 },
-        { key: "b", val: 2 },
-        { key: "a", val: 3 },
-        { key: "c", val: 4 },
+        { key: "a", visited: [] },
+        { key: "b", visited: [] },
+        { key: "a", visited: [] },
+        { key: "c", visited: [] },
     ];
-    const results = [
-        { key: "a", val: 2 },
-        { key: "b", val: 3 },
-        { key: "a", val: 4 },
-        { key: "c", val: 5 },
-    ];
-    const destinationStreamKeys = [];
-    const sink = new Writable({
-        objectMode: true,
-        write(chunk, enc, cb) {
-            expect(results).to.deep.include(chunk);
-            expect(input).to.not.deep.include(chunk);
-            t.pass();
-            cb();
-        },
-    });
-    const construct = (destKey: string) => {
-        destinationStreamKeys.push(destKey);
+
+    const construct = sinon.spy((destKey: string) => {
         const dest = map((chunk: Test) => {
-            return {
-                ...chunk,
-                val: chunk.val + 1,
-            };
+            chunk.visited.push(1);
+            return chunk;
         });
 
-        dest.pipe(sink);
         return dest;
-    };
+    });
 
     const demuxed = demux(
         construct,
-        { keyBy: (chunk: any) => chunk.key },
+        { keyBy: item => item.key },
         { objectMode: true },
     );
+
     demuxed.on("finish", () => {
-        expect(destinationStreamKeys).to.deep.equal(["a", "b", "c"]);
+        expect(construct.withArgs("a").callCount).to.equal(1);
+        expect(construct.withArgs("b").callCount).to.equal(1);
+        expect(construct.withArgs("c").callCount).to.equal(1);
         t.pass();
         t.end();
     });
@@ -110,17 +79,18 @@ test.cb("should spread per key using keyBy", t => {
 
 test.cb("should emit errors", t => {
     t.plan(2);
+    let index = 0;
     const input = [
-        { key: "a", val: 1 },
-        { key: "b", val: 2 },
-        { key: "a", val: 3 },
-        { key: "a", val: 4 },
+        { key: "a", visited: [] },
+        { key: "b", visited: [] },
+        { key: "a", visited: [] },
+        { key: "a", visited: [] },
     ];
     const results = [
-        { key: "a", val: 2 },
-        { key: "b", val: 3 },
-        { key: "a", val: 4 },
-        { key: "a", val: 5 },
+        { key: "a", visited: [0] },
+        { key: "b", visited: [1] },
+        { key: "a", visited: [2] },
+        { key: "a", visited: [3] },
     ];
     const destinationStreamKeys = [];
     const sink = new Writable({
@@ -131,7 +101,7 @@ test.cb("should emit errors", t => {
             t.pass();
             cb();
         },
-    }).on("unpipe", e => console.log("sink err"));
+    });
 
     const construct = (destKey: string) => {
         destinationStreamKeys.push(destKey);
@@ -139,11 +109,12 @@ test.cb("should emit errors", t => {
             if (chunk.key === "b") {
                 throw new Error("Caught object with key 'b'");
             }
-            return {
-                ...chunk,
-                val: chunk.val + 1,
-            };
-        }).on("error", e => console.log("got err"));
+
+            const _chunk = { ...chunk, visited: [] };
+            _chunk.visited.push(index);
+            index++;
+            return _chunk;
+        }).on("error", () => {});
 
         dest.pipe(sink);
         return dest;
@@ -162,3 +133,374 @@ test.cb("should emit errors", t => {
     input.forEach(event => demuxed.write(event));
     demuxed.end();
 });
+
+test("compose() should emit drain event ~rate * highWaterMark ms for every write that causes backpressure", async t => {
+    t.plan(7);
+    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) {
+                cb();
+                t.pass();
+                pendingReads--;
+                if (pendingReads === 0) {
+                    resolve();
+                }
+            },
+        });
+        const construct = (destKey: string) => {
+            const first = map(async (chunk: Chunk) => {
+                await sleep(_rate);
+                chunk.mapped.push(1);
+                return chunk;
+            });
+
+            const second = map(async (chunk: Chunk) => {
+                chunk.mapped.push(2);
+                return chunk;
+            });
+
+            first.pipe(second).pipe(sink);
+            return first;
+        };
+        const _demux = demux(
+            construct,
+            { key: "key" },
+            {
+                objectMode: true,
+                highWaterMark,
+            },
+        );
+        _demux.on("error", err => {
+            reject();
+        });
+
+        _demux.on("drain", () => {
+            expect(_demux._writableState.length).to.be.equal(0);
+            expect(performance.now() - start).to.be.greaterThan(_rate);
+            t.pass();
+        });
+
+        const input = [
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+        ];
+        let pendingReads = input.length;
+
+        let start = performance.now();
+        for (const item of input) {
+            const res = _demux.write(item);
+            expect(_demux._writableState.length).to.be.at.most(highWaterMark);
+            if (!res) {
+                start = performance.now();
+                await sleep(100);
+            }
+        }
+    });
+});
+
+test("demux() should emit one drain event when writing 6 items with highWaterMark of 5", t => {
+    t.plan(7);
+    const highWaterMark = 5;
+    return new Promise(async (resolve, reject) => {
+        interface Chunk {
+            key: string;
+            mapped: number[];
+        }
+        const sink = new Writable({
+            objectMode: true,
+            write(chunk, encoding, cb) {
+                cb();
+                t.pass();
+                if (chunk.key === "f") {
+                    resolve();
+                }
+            },
+        });
+        const construct = (destKey: string) => {
+            const first = map(async (chunk: Chunk) => {
+                chunk.mapped.push(1);
+                return chunk;
+            });
+
+            const second = map(async (chunk: Chunk) => {
+                chunk.mapped.push(2);
+                return chunk;
+            });
+
+            first.pipe(second).pipe(sink);
+            return first;
+        };
+        const _demux = demux(
+            construct,
+            { key: "key" },
+            {
+                objectMode: true,
+                highWaterMark,
+            },
+        );
+        _demux.on("error", err => {
+            reject();
+        });
+
+        _demux.on("drain", () => {
+            expect(_demux._writableState.length).to.be.equal(0);
+            t.pass();
+        });
+
+        const input = [
+            { key: "a", mapped: [] },
+            { key: "b", mapped: [] },
+            { key: "c", mapped: [] },
+            { key: "d", mapped: [] },
+            { key: "e", mapped: [] },
+            { key: "f", mapped: [] },
+        ];
+
+        for (const item of input) {
+            const res = _demux.write(item);
+            expect(_demux._writableState.length).to.be.at.most(highWaterMark);
+            if (!res) {
+                await sleep(10);
+            }
+        }
+    });
+});
+
+test.cb(
+    "demux() should emit drain event after 500 ms when writing 5 items that take 100ms to process with a highWaterMark of 5 ",
+    t => {
+        t.plan(6);
+        const _rate = 100;
+        const highWaterMark = 5;
+        interface Chunk {
+            key: string;
+            mapped: number[];
+        }
+        const sink = new Writable({
+            objectMode: true,
+            write(chunk, encoding, cb) {
+                t.pass();
+                cb();
+                if (pendingReads === 0) {
+                    t.end();
+                }
+            },
+        });
+        const construct = (destKey: string) => {
+            const first = map(
+                async (chunk: Chunk) => {
+                    chunk.mapped.push(1);
+                    await sleep(_rate);
+                    return chunk;
+                },
+                { objectMode: true },
+            );
+
+            const second = map(
+                (chunk: Chunk) => {
+                    pendingReads--;
+                    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 => {
+            t.end(err);
+        });
+
+        _demux.on("drain", () => {
+            expect(_demux._writableState.length).to.be.equal(0);
+            expect(performance.now() - start).to.be.greaterThan(
+                _rate * input.length,
+            );
+            t.pass();
+        });
+
+        const input = [
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+        ];
+
+        let pendingReads = input.length;
+        input.forEach(item => {
+            _demux.write(item);
+        });
+        const start = performance.now();
+    },
+);
+test.cb(
+    "demux() should emit drain event immediately when second stream is bottleneck",
+    t => {
+        t.plan(6);
+        const highWaterMark = 5;
+        interface Chunk {
+            key: string;
+            mapped: number[];
+        }
+        const sink = new Writable({
+            objectMode: true,
+            write(chunk, encoding, cb) {
+                t.pass();
+                cb();
+                if (pendingReads === 0) {
+                    t.end();
+                }
+            },
+        });
+        const construct = (destKey: string) => {
+            const first = map(
+                (chunk: Chunk) => {
+                    chunk.mapped.push(1);
+                    return chunk;
+                },
+                { objectMode: true },
+            );
+
+            const second = map(
+                async (chunk: Chunk) => {
+                    pendingReads--;
+                    await sleep(200);
+                    chunk.mapped.push(2);
+                    expect(second._writableState.length).to.be.equal(1);
+                    expect(first._readableState.length).to.equal(pendingReads);
+                    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 => {
+            t.end(err);
+        });
+
+        _demux.on("drain", () => {
+            expect(_demux._writableState.length).to.be.equal(0);
+            expect(performance.now() - start).to.be.lessThan(50);
+            t.pass();
+        });
+
+        const input = [
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+        ];
+
+        let pendingReads = input.length;
+        input.forEach(item => {
+            _demux.write(item);
+        });
+        const start = performance.now();
+    },
+);
+
+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;
+    return new Promise(async (resolve, reject) => {
+        interface Chunk {
+            key: string;
+            mapped: number[];
+        }
+        const sink = new Writable({
+            objectMode: true,
+            write(chunk, encoding, cb) {
+                t.pass();
+                cb();
+                if (pendingReads === 0) {
+                    resolve();
+                }
+            },
+        });
+        const construct = (destKey: string) => {
+            const first = map(
+                (chunk: Chunk) => {
+                    expect(first._readableState.length).to.be.at.most(2);
+                    chunk.mapped.push(1);
+                    return chunk;
+                },
+                { objectMode: 2, highWaterMark: 2 },
+            );
+
+            const second = map(
+                async (chunk: Chunk) => {
+                    chunk.mapped.push(2);
+                    expect(second._writableState.length).to.be.equal(1);
+                    await sleep(100);
+                    pendingReads--;
+                    return chunk;
+                },
+                { objectMode: 2, highWaterMark: 2 },
+            );
+
+            first.pipe(second).pipe(sink);
+            return first;
+        };
+        const _demux = demux(
+            construct,
+            { key: "key" },
+            {
+                objectMode: true,
+                highWaterMark,
+            },
+        );
+        _demux.on("error", err => {
+            reject();
+        });
+
+        _demux.on("drain", () => {
+            expect(_demux._writableState.length).to.be.equal(0);
+            t.pass();
+        });
+
+        const input = [
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+            { key: "a", mapped: [] },
+        ];
+        let pendingReads = input.length;
+
+        input.forEach(item => {
+            _demux.write(item);
+        });
+    });
+});