Add unit tests

Author: tohhsinpei

src/commands/database-set.ts

@@ -1,7 +1,6 @@
 import * as clc from "colorette";
 import * as fs from "fs";
 
-import { Client } from "../apiv2";
 import { Command } from "../command";
 import { Emulators } from "../emulator/types";
 import { FirebaseError } from "../error";
@@ -14,7 +13,7 @@ import { URL } from "url";
 import { logger } from "../logger";
 import { requireDatabaseInstance } from "../requireDatabaseInstance";
 import * as utils from "../utils";
-import { DatabaseChunkUploader } from "../databaseChunkUploader";
+import DatabaseImporter from "../database/import";
 
 export const command = new Command("database:set <path> [infile]")
   .description("store JSON data at the specified path via STDIN, arg, or file")
@@ -61,9 +60,9 @@ export const command = new Command("database:set <path> [infile]")
       utils.explainStdin();
     }
 
-    const uploader = new DatabaseChunkUploader(dbUrl, inputString);
+    const importer = new DatabaseImporter(dbUrl, inputString);
     try {
-      await uploader.upload(/* overwrite= */ true);
+      await importer.execute(/* overwrite= */ true);
     } catch (err: any) {
       logger.debug(err);
       throw new FirebaseError(`Unexpected error while setting data: ${err}`, { exit: 2 });

src/commands/database-update.ts

@@ -2,7 +2,6 @@ import { URL } from "url";
 import * as clc from "colorette";
 import * as fs from "fs";
 
-import { Client } from "../apiv2";
 import { Command } from "../command";
 import { Emulators } from "../emulator/types";
 import { FirebaseError } from "../error";
@@ -14,7 +13,7 @@ import { requirePermissions } from "../requirePermissions";
 import { logger } from "../logger";
 import { requireDatabaseInstance } from "../requireDatabaseInstance";
 import * as utils from "../utils";
-import { DatabaseChunkUploader } from "../databaseChunkUploader";
+import DatabaseImporter from "../database/import";
 
 export const command = new Command("database:update <path> [infile]")
   .description("update some of the keys for the defined path in your Firebase")
@@ -61,9 +60,9 @@ export const command = new Command("database:update <path> [infile]")
       utils.explainStdin();
     }
 
-    const uploader = new DatabaseChunkUploader(dbUrl, inputString);
+    const importer = new DatabaseImporter(dbUrl, inputString);
     try {
-      await uploader.upload(/* overwrite= */ false);
+      await importer.execute(/* overwrite= */ false);
     } catch (err: any) {
       throw new FirebaseError("Unexpected error while setting data");
     }

src/databaseChunkUploader.ts renamed to src/database/import.ts

@@ -1,11 +1,10 @@
 import { URL } from "url";
-import { Client } from "./apiv2";
-import * as utils from "./utils";
+import { Client } from "../apiv2";
 
 const MAX_CHUNK_SIZE = 1024 * 1024;
 
 type Data = {
-  json: any;
+  json: { [key: string]: any } | string | number | boolean;
   pathname: string;
 };
 
@@ -14,19 +13,28 @@ type ChunkedData = {
   size: number;
 };
 
-export class DatabaseChunkUploader {
+/**
+ * Imports JSON data to a given RTDB instance.
+ *
+ * The data is parsed and chunked into subtrees of ~1 MB, to be subsequently written in parallel.
+ */
+export default class DatabaseImporter {
+  chunks: Data[];
   private client: Client;
-  private chunks: Data[];
-
-  constructor(private dbUrl: URL, file: string) {
-    this.client = new Client({ urlPrefix: dbUrl.origin, auth: true });
 
+  constructor(private dbUrl: URL, file: string, private chunkSize = MAX_CHUNK_SIZE) {
     const data = { json: JSON.parse(file), pathname: dbUrl.pathname };
     const chunkedData = this.chunkData(data);
     this.chunks = chunkedData.chunks || [data];
+    this.client = new Client({ urlPrefix: dbUrl.origin, auth: true });
   }
 
-  public async upload(overwrite: boolean): Promise<any> {
+  /**
+   * Writes the chunked data to RTDB.
+   *
+   * @param overwrite Whether to overwrite the existing data at the given location.
+   */
+  async execute(overwrite: boolean): Promise<any> {
     return Promise.all(
       this.chunks.map((chunk: Data) =>
         this.client.request({
@@ -40,16 +48,20 @@ export class DatabaseChunkUploader {
   }
 
   private chunkData({ json, pathname }: Data): ChunkedData {
-    if (isObject(json)) {
+    if (typeof json === "string" || typeof json === "number" || typeof json === "boolean") {
+      // Leaf node, cannot be chunked
+      return { chunks: null, size: JSON.stringify(json).length };
+    } else {
       // Children node
       let size = 2; // {}
-      let chunks = [];
+
+      const chunks = [];
       let hasChunkedChild = false;
 
-      for (const key in json) {
+      for (const key of Object.keys(json)) {
         size += key.length + 3; // "[key]":
 
-        const child = { json: json[key], pathname: pathname + "/" + key };
+        const child = { json: json[key], pathname: [pathname, key].join("/").replace("//", "/") };
         const childChunks = this.chunkData(child);
         size += childChunks.size;
         if (childChunks.chunks) {
@@ -60,18 +72,11 @@ export class DatabaseChunkUploader {
         }
       }
 
-      if (hasChunkedChild || size >= MAX_CHUNK_SIZE) {
+      if (hasChunkedChild || size >= this.chunkSize) {
         return { chunks, size };
       } else {
         return { chunks: null, size };
       }
-    } else {
-      // Leaf node, cannot be chunked
-      return { chunks: null, size: JSON.stringify(json).length };
     }
   }
 }
-
-function isObject(blob: any): boolean {
-  return blob !== null && typeof blob === "object";
-}

src/test/database/import.spec.ts

@@ -0,0 +1,25 @@
+import { expect } from "chai";
+
+import DatabaseImporter from "../../database/import";
+
+const dbUrl = new URL("https://test-db.firebaseio.com/foo");
+
+describe("DatabaseImporter", () => {
+  const DATA = { a: 100, b: { c: true, d: { e: "bar", f: { g: 0, h: 1 } } } };
+
+  it("parses data as single chunk", () => {
+    const importer = new DatabaseImporter(dbUrl, JSON.stringify(DATA));
+    expect(importer.chunks.length).to.equal(1);
+    expect(importer.chunks[0].json).to.deep.equal(DATA);
+    expect(importer.chunks[0].pathname).to.equal("/foo");
+  });
+
+  it("parses data as multiple chunks", () => {
+    const importer = new DatabaseImporter(dbUrl, JSON.stringify(DATA), /* chunkSize= */ 20);
+    expect(importer.chunks.length).to.equal(4);
+    expect(importer.chunks).to.deep.include({ json: 100, pathname: "/foo/a" });
+    expect(importer.chunks).to.deep.include({ json: true, pathname: "/foo/b/c" });
+    expect(importer.chunks).to.deep.include({ json: "bar", pathname: "/foo/b/d/e" });
+    expect(importer.chunks).to.deep.include({ json: { g: 0, h: 1 }, pathname: "/foo/b/d/f" });
+  });
+});