refactor existing code

unipept · Aug 5, 2024 · 3b575a9 · 3b575a9
1 parent c9aa3aa
commit 3b575a9
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 22 deletions.
diff --git a/lib/commands/unipept/unipept_subcommand.ts b/lib/commands/unipept/unipept_subcommand.ts
@@ -20,13 +20,17 @@ export abstract class UnipeptSubcommand {
   selectedFields?: RegExp[];
   fasta: boolean;
 
+  // we must save this to be able to close it properly in tests
+  private streamInterface?: Interface;
+
   constructor(name: string) {
     this.name = name;
     const version = JSON.parse(readFileSync(new URL("../../../package.json", import.meta.url), "utf8")).version;
     this.user_agent = `unipept-cli/${version}`;
     this.command = this.create(name);
     this.fasta = false;
   }
+
   abstract defaultBatchSize(): number;
 
   requiredFields(): string[] {
@@ -58,16 +62,10 @@ export abstract class UnipeptSubcommand {
       this.outputStream = createWriteStream(this.options.output);
     }
 
-    let slice = [];
+    const iterator = this.getInputIterator(args, options.input);
+    const firstLine = (await iterator.next()).value;
 
-    for await (const input of this.getInputIterator(args, options.input)) {
-      slice.push(input);
-      if (slice.length >= this.batchSize) {
-        await this.processBatch(slice);
-        slice = [];
-      }
-    }
-    await this.processBatch(slice);
+    await this.normalInputProcessor(firstLine, iterator);
   }
 
   async processBatch(slice: string[]): Promise<void> {
@@ -92,6 +90,19 @@ export abstract class UnipeptSubcommand {
     if (this.firstBatch) this.firstBatch = false;
   }
 
+  async normalInputProcessor(firstLine: string, iterator: IterableIterator<string> | AsyncIterableIterator<string>) {
+    let slice = [firstLine];
+
+    for await (const line of iterator) {
+      slice.push(line);
+      if (slice.length >= this.batchSize) {
+        await this.processBatch(slice);
+        slice = [];
+      }
+    }
+    await this.processBatch(slice);
+  }
+
   private constructRequestBody(slice: string[]): URLSearchParams {
     const names = this.getSelectedFields().length === 0 || this.getSelectedFields().some(regex => regex.toString().includes("name") || regex.toString().includes(".*$"));
     return new URLSearchParams({
@@ -128,13 +139,15 @@ export abstract class UnipeptSubcommand {
    * - if an input file is given, use the file
    * - otherwise, use standard input
    */
-  private getInputIterator(args: string[], input?: string): string[] | Interface {
+  private getInputIterator(args: string[], input?: string): IterableIterator<string> | AsyncIterableIterator<string> {
     if (args.length > 0) {
-      return args;
+      return args.values();
     } else if (input) {
-      return createInterface({ input: createReadStream(input) });
+      this.streamInterface = createInterface({ input: createReadStream(input) });
+      return this.streamInterface[Symbol.asyncIterator]();
     } else {
-      return createInterface({ input: process.stdin })
+      this.streamInterface = createInterface({ input: process.stdin });
+      return this.streamInterface[Symbol.asyncIterator]();
     }
   }
 

diff --git a/tests/commands/unipept/unipept_subcommand.test.ts b/tests/commands/unipept/unipept_subcommand.test.ts
@@ -28,19 +28,18 @@ test('test correct inputIterator', async () => {
   const command = new Pept2lca();
 
   // should be stdin
-  let input = command["getInputIterator"]([]) as Interface;
-  expect(input).toBeInstanceOf(Interface);
-  input.close();
+  let input = command["getInputIterator"]([]) as AsyncIterableIterator<string>;
+  expect(typeof input[Symbol.asyncIterator]).toBe("function");
+  command['streamInterface']?.close();
 
   // should be a (non-existant) file and error
-  input = command["getInputIterator"]([], "filename") as Interface;
-  input.on("error", (e) => {
-    expect(e.toString()).toMatch(/no such file/);
-  });
+  input = command["getInputIterator"]([], "filename") as AsyncIterableIterator<string>;
+  expect(typeof input[Symbol.asyncIterator]).toBe("function");
+  await expect(async () => { await input.next() }).rejects.toThrow(/no such file/);
 
   // should be array
-  const inputArray = command["getInputIterator"](["A", "B"]);
-  expect(inputArray).toBeInstanceOf(Array);
+  const inputArray = command["getInputIterator"](["A", "B"]) as IterableIterator<string>;
+  expect(typeof inputArray[Symbol.iterator]).toBe("function");
 });
 
 test('test selected fields parsing', () => {