improve bindings test infra

duckdb · Oct 5, 2024 · d756090 · d756090
1 parent 978fc97
commit d756090
Show file tree

Hide file tree

Showing 18 changed files with 1,152 additions and 436 deletions.
diff --git a/bindings/test/prepared_statements.test.ts b/bindings/test/prepared_statements.test.ts
@@ -1,27 +1,28 @@
 import duckdb from '@duckdb/node-bindings';
 import { expect, suite, test } from 'vitest';
+import { data } from './utils/expectedVectors';
 import { expectResult } from './utils/expectResult';
 import { withConnection } from './utils/withConnection';
 
 suite('prepared statements', () => {
   test('no parameters', async () => {
-    await withConnection(async (con) => {
-      const prepared = await duckdb.prepare(con, 'select 17 as seventeen');
+    await withConnection(async (connection) => {
+      const prepared = await duckdb.prepare(connection, 'select 17 as seventeen');
       try {
         expect(duckdb.nparams(prepared)).toBe(0);
         expect(duckdb.prepared_statement_type(prepared)).toBe(duckdb.StatementType.SELECT);
-        const res = await duckdb.execute_prepared(prepared);
+        const result = await duckdb.execute_prepared(prepared);
         try {
-          await expectResult(res, {
+          await expectResult(result, {
             columns: [
-              { name: 'seventeen', type: duckdb.Type.INTEGER },
+              { name: 'seventeen', logicalType: { typeId: duckdb.Type.INTEGER } },
             ],
             chunks: [
-              { rowCount: 1, vectors: [{ byteCount: 4, validity: [true], values: [17] }]},
+              { rowCount: 1, vectors: [data(4, [true], [17])]},
             ],
           });
         } finally {
-          duckdb.destroy_result(res);
+          duckdb.destroy_result(result);
         }
       } finally {
         duckdb.destroy_prepare(prepared);

diff --git a/bindings/test/query.test.ts b/bindings/test/query.test.ts
diff --git a/bindings/test/utils/ExpectedLogicalType.ts b/bindings/test/utils/ExpectedLogicalType.ts
@@ -0,0 +1,74 @@
+import duckdb from '@duckdb/node-bindings';
+
+export interface ExpectedSimpleLogicalType {
+  typeId: Exclude<duckdb.Type,
+    | duckdb.Type.ARRAY
+    | duckdb.Type.DECIMAL
+    | duckdb.Type.ENUM
+    | duckdb.Type.LIST
+    | duckdb.Type.MAP
+    | duckdb.Type.STRUCT
+    | duckdb.Type.UNION
+  >;
+}
+
+export interface ExpectedArrayLogicalType {
+  typeId: duckdb.Type.ARRAY;
+  valueType: ExpectedLogicalType;
+  size: number;
+}
+
+export interface ExpectedDecimalLogicalType {
+  typeId: duckdb.Type.DECIMAL;
+  width: number;
+  scale: number;
+  internalType: duckdb.Type;
+}
+
+export interface ExpectedEnumLogicalType {
+  typeId: duckdb.Type.ENUM;
+  values: string[];
+  internalType: duckdb.Type;
+}
+
+export interface ExpectedListLogicalType {
+  typeId: duckdb.Type.LIST;
+  valueType: ExpectedLogicalType;
+}
+
+export interface ExpectedMapLogicalType {
+  typeId: duckdb.Type.MAP;
+  keyType: ExpectedLogicalType;
+  valueType: ExpectedLogicalType;
+}
+
+export interface ExpectedStructEntry {
+  name: string;
+  type: ExpectedLogicalType;
+}
+
+export interface ExpectedStructLogicalType {
+  typeId: duckdb.Type.STRUCT;
+  entries: ExpectedStructEntry[];
+}
+
+export interface ExpectedUnionAlternative {
+  tag: string;
+  type: ExpectedLogicalType;
+}
+
+export interface ExpectedUnionLogicalType {
+  typeId: duckdb.Type.UNION;
+  alternatives: ExpectedUnionAlternative[];
+}
+
+export type ExpectedLogicalType =
+  | ExpectedSimpleLogicalType
+  | ExpectedArrayLogicalType
+  | ExpectedDecimalLogicalType
+  | ExpectedEnumLogicalType
+  | ExpectedListLogicalType
+  | ExpectedMapLogicalType
+  | ExpectedStructLogicalType
+  | ExpectedUnionLogicalType
+  ;
diff --git a/bindings/test/utils/ExpectedResult.ts b/bindings/test/utils/ExpectedResult.ts
@@ -0,0 +1,22 @@
+import duckdb from '@duckdb/node-bindings';
+import { ExpectedLogicalType } from './ExpectedLogicalType';
+import { ExpectedVector } from './ExpectedVector';
+
+export interface ExpectedColumn {
+  name: string;
+  logicalType: ExpectedLogicalType;
+}
+
+export interface ExpectedChunk {
+  columnCount?: number;
+  rowCount: number;
+  vectors: ExpectedVector[];
+}
+
+export interface ExpectedResult {
+  statementType?: duckdb.StatementType;
+  resultType?: duckdb.ResultType;
+  rowsChanged?: number;
+  columns: ExpectedColumn[];
+  chunks: ExpectedChunk[];
+}
diff --git a/bindings/test/utils/ExpectedVector.ts b/bindings/test/utils/ExpectedVector.ts
@@ -0,0 +1,52 @@
+export interface ExpectedArrayVector {
+  kind: 'array';
+  itemCount: number;
+  validity: boolean[];
+  child: ExpectedVector;
+}
+
+export interface ExpectedDataVector {
+  kind: 'data';
+  validity: boolean[];
+  itemBytes: number;
+  values: any[];
+}
+
+export type ExpectedListEntry = [bigint, bigint] | null;
+
+export interface ExpectedListVector {
+  kind: 'list';
+  validity: boolean[];
+  entries: (ExpectedListEntry | null)[];
+  childItemCount: number;
+  child: ExpectedVector;
+}
+
+export interface ExpectedMapVector {
+  kind: 'map';
+  validity: boolean[];
+  entries: (ExpectedListEntry | null)[];
+  keys: ExpectedVector;
+  values: ExpectedVector;
+}
+
+export interface ExpectedStructVector {
+  kind: 'struct';
+  itemCount: number;
+  validity: boolean[];
+  children: ExpectedVector[];
+}
+
+export interface ExpectedUnionVector {
+  kind: 'union';
+  children: ExpectedVector[];
+}
+
+export type ExpectedVector =
+  | ExpectedArrayVector
+  | ExpectedDataVector
+  | ExpectedListVector
+  | ExpectedMapVector
+  | ExpectedStructVector
+  | ExpectedUnionVector
+  ;
diff --git a/bindings/test/utils/expectChunk.ts b/bindings/test/utils/expectChunk.ts
@@ -0,0 +1,23 @@
+import duckdb from '@duckdb/node-bindings';
+import { expect } from 'vitest';
+import { ExpectedChunk, ExpectedColumn } from './ExpectedResult';
+import { expectLogicalType } from './expectLogicalType';
+import { expectVector } from './expectVector';
+import { withLogicalType } from './withLogicalType';
+
+export function expectChunk(chunk: duckdb.DataChunk, expectedChunk: ExpectedChunk, expectedColumns: ExpectedColumn[]) {
+  const chunkColumnCount = expectedChunk.columnCount ?? expectedColumns.length;
+  expect(duckdb.data_chunk_get_column_count(chunk)).toBe(chunkColumnCount);
+  expect(duckdb.data_chunk_get_size(chunk)).toBe(expectedChunk.rowCount);
+  for (let col = 0; col < expectedChunk.vectors.length; col++) {
+    const expectedVector = expectedChunk.vectors[col];
+    const vector = duckdb.data_chunk_get_vector(chunk, col);
+
+    const expectedLogicalType = expectedColumns[col].logicalType;
+    withLogicalType(duckdb.vector_get_column_type(vector),
+      (logical_type) => expectLogicalType(logical_type, expectedLogicalType, `col ${col}`)
+    );
+
+    expectVector(vector, expectedVector, expectedLogicalType, `col ${col}`);
+  }
+}
diff --git a/bindings/test/utils/expectLogicalType.ts b/bindings/test/utils/expectLogicalType.ts
@@ -0,0 +1,48 @@
+import duckdb from '@duckdb/node-bindings';
+import { expect } from 'vitest';
+import { ExpectedLogicalType } from './ExpectedLogicalType';
+
+export function expectLogicalType(logical_type: duckdb.LogicalType, expectedLogicalType: ExpectedLogicalType, message?: string) {
+  expect(duckdb.get_type_id(logical_type), message).toBe(expectedLogicalType.typeId);
+  switch (expectedLogicalType.typeId) {
+    case duckdb.Type.ARRAY:
+      expectLogicalType(duckdb.array_type_child_type(logical_type), expectedLogicalType.valueType);
+      expect(duckdb.array_type_array_size(logical_type)).toBe(expectedLogicalType.size);
+      break;
+    case duckdb.Type.DECIMAL:
+      expect(duckdb.decimal_width(logical_type)).toBe(expectedLogicalType.width);
+      expect(duckdb.decimal_scale(logical_type)).toBe(expectedLogicalType.scale);
+      expect(duckdb.decimal_internal_type(logical_type)).toBe(expectedLogicalType.internalType);
+      break;
+    case duckdb.Type.ENUM:
+      {
+        expect(duckdb.enum_internal_type(logical_type)).toBe(expectedLogicalType.internalType);
+        expect(duckdb.enum_dictionary_size(logical_type)).toBe(expectedLogicalType.values.length);
+        for (let i = 0; i < expectedLogicalType.values.length; i++) {
+          expect(duckdb.enum_dictionary_value(logical_type, i)).toBe(expectedLogicalType.values[i]);
+        }
+      }
+      break;
+    case duckdb.Type.LIST:
+      expectLogicalType(duckdb.list_type_child_type(logical_type), expectedLogicalType.valueType);
+      break;
+    case duckdb.Type.MAP:
+      expectLogicalType(duckdb.map_type_key_type(logical_type), expectedLogicalType.keyType);
+      expectLogicalType(duckdb.map_type_value_type(logical_type), expectedLogicalType.valueType);
+      break;
+    case duckdb.Type.STRUCT:
+      expect(duckdb.struct_type_child_count(logical_type)).toBe(expectedLogicalType.entries.length);
+      for (let i = 0; i < expectedLogicalType.entries.length; i++) {
+        expect(duckdb.struct_type_child_name(logical_type, i)).toBe(expectedLogicalType.entries[i].name);
+        expectLogicalType(duckdb.struct_type_child_type(logical_type, i), expectedLogicalType.entries[i].type);
+      }
+      break;
+    case duckdb.Type.UNION:
+      expect(duckdb.union_type_member_count(logical_type)).toBe(expectedLogicalType.alternatives.length);
+      for (let i = 0; i < expectedLogicalType.alternatives.length; i++) {
+        expect(duckdb.union_type_member_name(logical_type, i)).toBe(expectedLogicalType.alternatives[i].tag);
+        expectLogicalType(duckdb.union_type_member_type(logical_type, i), expectedLogicalType.alternatives[i].type);
+      }
+      break;
+  }
+}
diff --git a/bindings/test/utils/expectResult.ts b/bindings/test/utils/expectResult.ts
@@ -1,74 +1,27 @@
 import duckdb from '@duckdb/node-bindings';
 import { expect } from 'vitest';
-import { expectValidity } from './validityTestUtils';
-import { getValue } from './valueTestUtils';
+import { expectChunk } from './expectChunk';
+import { ExpectedResult } from './ExpectedResult';
+import { expectLogicalType } from './expectLogicalType';
+import { withLogicalType } from './withLogicalType';
 
-export interface ExpectedResult {
-  statementType?: duckdb.StatementType;
-  resultType?: duckdb.ResultType;
-  rowsChanged?: number;
-  columns: ExpectedColumn[];
-  chunks: ExpectedChunk[];
-}
-
-export interface ExpectedColumn {
-  name: string;
-  type: duckdb.Type;
-}
-
-export interface ExpectedChunk {
-  rowCount: number;
-  vectors: ExpectedVector[];
-}
-
-export interface ExpectedVector {
-  byteCount: number;
-  validity: boolean[];
-  values: any[];
-}
-
-export async function expectResult(res: duckdb.Result, expected: ExpectedResult) {
-  expect(duckdb.result_statement_type(res)).toBe(expected.statementType ?? duckdb.StatementType.SELECT);
-  expect(duckdb.result_return_type(res)).toBe(expected.resultType ?? duckdb.ResultType.QUERY_RESULT);
-  expect(duckdb.rows_changed(res)).toBe(expected.rowsChanged ?? 0);
-  expect(duckdb.column_count(res)).toBe(expected.columns.length);
-  for (let col = 0; col < expected.columns.length; col++) {
-    const expectedColumn = expected.columns[col];
-    expect(duckdb.column_name(res, col)).toBe(expectedColumn.name);
-    expect(duckdb.column_type(res, col)).toBe(expectedColumn.type);
-    const logical_type = duckdb.column_logical_type(res, col);
-    try {
-      expect(duckdb.get_type_id(logical_type)).toBe(expectedColumn.type);
-    } finally {
-      duckdb.destroy_logical_type(logical_type);
-    }
+export async function expectResult(result: duckdb.Result, expectedResult: ExpectedResult) {
+  expect(duckdb.result_statement_type(result)).toBe(expectedResult.statementType ?? duckdb.StatementType.SELECT);
+  expect(duckdb.result_return_type(result)).toBe(expectedResult.resultType ?? duckdb.ResultType.QUERY_RESULT);
+  expect(duckdb.rows_changed(result)).toBe(expectedResult.rowsChanged ?? 0);
+  expect(duckdb.column_count(result)).toBe(expectedResult.columns.length);
+  for (let col = 0; col < expectedResult.columns.length; col++) {
+    const expectedColumn = expectedResult.columns[col];
+    expect(duckdb.column_name(result, col)).toBe(expectedColumn.name);
+    expect(duckdb.column_type(result, col)).toBe(expectedColumn.logicalType.typeId);
+    withLogicalType(duckdb.column_logical_type(result, col),
+      (logical_type) => expectLogicalType(logical_type, expectedColumn.logicalType, `col ${col}`)
+    );
   }
-  for (const expectedChunk of expected.chunks) {
-    const chunk = await duckdb.fetch_chunk(res);
+  for (const expectedChunk of expectedResult.chunks) {
+    const chunk = await duckdb.fetch_chunk(result);
     try {
-      expect(duckdb.data_chunk_get_column_count(chunk)).toBe(expected.columns.length);
-      expect(duckdb.data_chunk_get_size(chunk)).toBe(expectedChunk.rowCount);
-      for (let col = 0; col < expected.columns.length; col++) {
-        const expectedColumn = expected.columns[col];
-        const expectedVector = expectedChunk.vectors[col];
-        const vector = duckdb.data_chunk_get_vector(chunk, col);
-        const logical_type = duckdb.vector_get_column_type(vector);
-        try {
-          expect(duckdb.get_type_id(logical_type)).toBe(expectedColumn.type);
-        } finally {
-          duckdb.destroy_logical_type(logical_type);
-        }
-        const uint64Count = Math.ceil(expectedChunk.rowCount / 64);
-        const byteCount = uint64Count * 8;
-        const validity_bytes = duckdb.vector_get_validity(vector, byteCount);
-        const validity = new BigUint64Array(validity_bytes.buffer, 0, uint64Count);
-        const data = duckdb.vector_get_data(vector, expectedVector.byteCount);
-        const dv = new DataView(data.buffer);
-        for (let row = 0; row < expectedChunk.rowCount; row++) {
-          expectValidity(validity_bytes, validity, row, expectedVector.validity[row]);
-          expect(getValue(expectedColumn.type, dv, row)).toBe(expectedVector.values[row]);
-        }
-      }
+      expectChunk(chunk, expectedChunk, expectedResult.columns);
     } finally {
       duckdb.destroy_data_chunk(chunk);
     }

diff --git a/bindings/test/utils/expectValidity.ts b/bindings/test/utils/expectValidity.ts
@@ -0,0 +1,8 @@
+import duckdb from '@duckdb/node-bindings';
+import { expect } from 'vitest';
+import { isValid } from './isValid';
+
+export function expectValidity(validity_bytes: Uint8Array, validity: BigUint64Array, bit: number, expected: boolean, vectorName: string) {
+  expect(duckdb.validity_row_is_valid(validity_bytes, bit), `${vectorName} validity_bytes_bit[${bit}]`).toBe(expected);
+  expect(isValid(validity, bit), `${vectorName} validity_bit[${bit}]`).toBe(expected);
+}