1 files changed, 162 insertions, 0 deletions
diff --git a/src/arrow/js/test/jest-extensions.ts b/src/arrow/js/test/jest-extensions.ts
new file mode 100644
index 000000000..6adde0b83
--- /dev/null
+++ b/src/arrow/js/test/jest-extensions.ts
@@ -0,0 +1,162 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import { zip } from 'ix/iterable/zip';
+import { Table, Vector, RecordBatch, Column, util } from './Arrow';
+
+declare global {
+    namespace jest {
+        interface Matchers<R> {
+            toArrowCompare(expected: any): CustomMatcherResult;
+            toEqualTable(expected: Table): CustomMatcherResult;
+            toEqualRecordBatch(expected: RecordBatch): CustomMatcherResult;
+            toEqualVector(expected: Vector | [Vector | null, string?, string?]): CustomMatcherResult;
+        }
+    }
+}
+
+expect.extend({
+    toEqualTable,
+    toEqualVector,
+    toArrowCompare,
+    toEqualRecordBatch
+});
+
+function format(jest: jest.MatcherUtils, actual: any, expected: any, msg= ' ') {
+    return `${
+        jest.utils.printReceived(actual)
+        }${msg}${
+        jest.utils.printExpected(expected)
+    }`;
+}
+
+function toArrowCompare(this: jest.MatcherUtils, actual: any, expected: any) {
+    if (!util.createElementComparator(expected)(actual)) {
+        return { pass: false, message: () => format(this, actual, expected, ' should equal ') };
+    }
+    return { pass: true, message: () => '' };
+}
+
+function toEqualTable(this: jest.MatcherUtils, actual: Table, expected: Table) {
+    const failures = [] as string[];
+    try { expect(actual).toHaveLength(expected.length); } catch (e) { failures.push(`${e}`); }
+    try { expect(actual.numCols).toEqual(expected.numCols); } catch (e) { failures.push(`${e}`); }
+    try { expect(actual.schema.metadata).toEqual(expected.schema.metadata); } catch (e) { failures.push(`${e}`); }
+    (() => {
+        for (let i = -1, n = actual.numCols; ++i < n;) {
+            const v1 = actual.getColumnAt(i);
+            const v2 = expected.getColumnAt(i);
+            const name = actual.schema.fields[i].name;
+            try {
+                expect([v1, `actual`, name]).toEqualVector([v2, `expected`, name]);
+            } catch (e) { failures.push(`${e}`); }
+        }
+    })();
+    return {
+        pass: failures.length === 0,
+        message: () => failures.join('\n'),
+    };
+}
+
+function toEqualRecordBatch(this: jest.MatcherUtils, actual: RecordBatch, expected: RecordBatch) {
+    const failures = [] as string[];
+    try { expect(actual).toHaveLength(expected.length); } catch (e) { failures.push(`${e}`); }
+    try { expect(actual.numCols).toEqual(expected.numCols); } catch (e) { failures.push(`${e}`); }
+    (() => {
+        for (let i = -1, n = actual.numCols; ++i < n;) {
+            const v1 = actual.getChildAt(i);
+            const v2 = expected.getChildAt(i);
+            const name = actual.schema.fields[i].name;
+            try {
+                expect([v1, `actual`, name]).toEqualVector([v2, `expected`, name]);
+            } catch (e) { failures.push(`${e}`); }
+        }
+    })();
+    return {
+        pass: failures.length === 0,
+        message: () => failures.join('\n'),
+    };
+}
+
+function toEqualVector<
+    TActual extends Vector | [Vector | null, string?, string?],
+    TExpected extends Vector | [Vector | null, string?]
+>(this: jest.MatcherUtils, actual: TActual, expected: TExpected) {
+
+    let [v1, format1 = '', columnName = ''] = Array.isArray(actual) ? actual : [actual];
+    let [v2, format2 = ''] = Array.isArray(expected) ? expected : [expected];
+
+    if (v1 instanceof Column && columnName === '') { columnName = v1.name; }
+
+    if (v1 == null || v2 == null) {
+        return {
+            pass: false,
+            message: () => [
+                [columnName, `(${format(this, format1, format2, ' !== ')})`].filter(Boolean).join(':'),
+                `${v1 == null ? 'actual' : 'expected'} is null`
+            ].join('\n')
+        };
+    }
+
+    let getFailures = new Array<string>();
+    let propsFailures = new Array<string>();
+    let iteratorFailures = new Array<string>();
+    let allFailures = [
+        { title: 'get', failures: getFailures },
+        { title: 'props', failures: propsFailures },
+        { title: 'iterator', failures: iteratorFailures }
+    ];
+
+    let props: (keyof Vector)[] = ['type', 'length', 'nullCount'];
+
+    (() => {
+        for (let i = -1, n = props.length; ++i < n;) {
+            const prop = props[i];
+            if (`${v1[prop]}` !== `${v2[prop]}`) {
+                propsFailures.push(`${prop}: ${format(this, v1[prop], v2[prop], ' !== ')}`);
+            }
+        }
+    })();
+
+    (() => {
+        for (let i = -1, n = v1.length; ++i < n;) {
+            let x1 = v1.get(i), x2 = v2.get(i);
+            if (!util.createElementComparator(x2)(x1)) {
+                getFailures.push(`${i}: ${format(this, x1, x2, ' !== ')}`);
+            }
+        }
+    })();
+
+    (() => {
+        let i = -1;
+        for (let [x1, x2] of zip(v1, v2)) {
+            ++i;
+            if (!util.createElementComparator(x2)(x1)) {
+                iteratorFailures.push(`${i}: ${format(this, x1, x2, ' !== ')}`);
+            }
+        }
+    })();
+
+    return {
+        pass: allFailures.every(({ failures }) => failures.length === 0),
+        message: () => [
+            [columnName, `(${format(this, format1, format2, ' !== ')})`].filter(Boolean).join(':'),
+            ...allFailures.map(({ failures, title }) =>
+                !failures.length ? `` : [`${title}:`, ...failures].join(`\n`))
+        ].join('\n')
+    };
+}