Refactor exportCsvResults and create test

1. `exportCsvResults` now no longer requires an `onFinish` callback.
2. The test adds a generic framework for creating a mock cli server.
   This should be used in future tests.

Author: aeisenberg

extensions/ql-vscode/src/query-history.ts

@@ -956,11 +956,11 @@ export class QueryHistoryManager extends DisposableObject {
       void this.tryOpenExternalFile(query.csvPath);
       return;
     }
-    await query.exportCsvResults(this.qs, query.csvPath, () => {
+    if (await query.exportCsvResults(this.qs, query.csvPath)) {
       void this.tryOpenExternalFile(
         query.csvPath
       );
-    });
+    }
   }
 
   async handleViewCsvAlerts(

extensions/ql-vscode/src/run-queries.ts

@@ -341,24 +341,33 @@ export class QueryEvaluationInfo {
   /**
    * Creates the CSV file containing the results of this query. This will only be called if the query
    * does not have interpreted results and the CSV file does not already exist.
+   *
+   * @return Promise<true> if the operation creates the file. Promise<false> if the operation does
+   * not create the file.
+   *
+   * @throws Error if the operation fails.
    */
-  async exportCsvResults(qs: qsClient.QueryServerClient, csvPath: string, onFinish: () => void): Promise<void> {
-    let stopDecoding = false;
-    const out = fs.createWriteStream(csvPath);
-    out.on('finish', onFinish);
-    out.on('error', () => {
-      if (!stopDecoding) {
-        stopDecoding = true;
-        void showAndLogErrorMessage(`Failed to write CSV results to ${csvPath}`);
-      }
-    });
+  async exportCsvResults(qs: qsClient.QueryServerClient, csvPath: string): Promise<boolean> {
     const resultSet = await this.chooseResultSet(qs);
     if (!resultSet) {
       void showAndLogWarningMessage('Query has no result set.');
-      return;
+      return false;
     }
+    let stopDecoding = false;
+    const out = fs.createWriteStream(csvPath);
+
+    const promise: Promise<boolean> = new Promise((resolve, reject) => {
+      out.on('finish', () => resolve(true));
+      out.on('error', () => {
+        if (!stopDecoding) {
+          stopDecoding = true;
+          reject(new Error(`Failed to write CSV results to ${csvPath}`));
+        }
+      });
+    });
+
     let nextOffset: number | undefined = 0;
-    while (nextOffset !== undefined && !stopDecoding) {
+    do {
       const chunk: DecodedBqrsChunk = await qs.cliServer.bqrsDecode(this.resultsPaths.resultsPath, resultSet, {
         pageSize: 100,
         offset: nextOffset,
@@ -370,8 +379,10 @@ export class QueryEvaluationInfo {
         }).join(',') + '\n');
       });
       nextOffset = chunk.next;
-    }
+    } while (nextOffset && !stopDecoding);
     out.end();
+
+    return promise;
   }
 
   /**

extensions/ql-vscode/src/vscode-tests/no-workspace/run-queries.test.ts

@@ -1,11 +1,16 @@
 import { expect } from 'chai';
 import * as path from 'path';
+import * as fs from 'fs-extra';
 import * as sinon from 'sinon';
 import { Uri } from 'vscode';
 
 import { QueryEvaluationInfo } from '../../run-queries';
 import { Severity, compileQuery } from '../../pure/messages';
 import * as config from '../../config';
+import { tmpDir } from '../../helpers';
+import { QueryServerClient } from '../../queryserver-client';
+import { CodeQLCliServer } from '../../cli';
+import { SELECT_QUERY_NAME } from '../../contextual/locationFinder';
 
 describe('run-queries', () => {
   let sandbox: sinon.SinonSandbox;
@@ -53,6 +58,51 @@ describe('run-queries', () => {
     expect(info.canHaveInterpretedResults()).to.eq(true);
   });
 
+  [SELECT_QUERY_NAME, 'other'].forEach(resultSetName => {
+    it(`should export csv results for result set ${resultSetName}`, async () => {
+      const csvLocation = path.join(tmpDir.name, 'test.csv');
+      const qs = createMockQueryServerClient(
+        createMockCliServer({
+          bqrsInfo: [{ 'result-sets': [{ name: resultSetName }, { name: 'hucairz' }] }],
+          bqrsDecode: [{
+            columns: [{ kind: 'NotString' }, { kind: 'String' }],
+            tuples: [['a', 'b'], ['c', 'd']],
+            next: 1
+          }, {
+            // just for fun, give a different set of columns here
+            // this won't happen with the real CLI, but it's a good test
+            columns: [{ kind: 'String' }, { kind: 'NotString' }, { kind: 'StillNotString' }],
+            tuples: [['a', 'b', 'c']]
+          }]
+        })
+      );
+      const info = createMockQueryInfo();
+      const promise = info.exportCsvResults(qs, csvLocation);
+
+      const result = await promise;
+      expect(result).to.eq(true);
+
+      const csv = fs.readFileSync(csvLocation, 'utf8');
+      expect(csv).to.eq('a,"b"\nc,"d"\n"a",b,c\n');
+
+      // now verify that we are using the expected result set
+      expect((qs.cliServer.bqrsDecode as sinon.SinonStub).callCount).to.eq(2);
+      expect((qs.cliServer.bqrsDecode as sinon.SinonStub).getCall(0).args[1]).to.eq(resultSetName);
+    });
+  });
+
+  it('should handle csv exports for a query with no result sets', async () => {
+    const csvLocation = path.join(tmpDir.name, 'test.csv');
+    const qs = createMockQueryServerClient(
+      createMockCliServer({
+        bqrsInfo: [{ 'result-sets': [] }]
+      })
+    );
+    const info = createMockQueryInfo();
+    const result = await info.exportCsvResults(qs, csvLocation);
+    expect(result).to.eq(false);
+  });
+
   describe('compile', () => {
     it('should compile', async () => {
       const info = createMockQueryInfo();
@@ -116,7 +166,7 @@ describe('run-queries', () => {
     );
   }
 
-  function createMockQueryServerClient() {
+  function createMockQueryServerClient(cliServer?: CodeQLCliServer): QueryServerClient {
     return {
       config: {
         timeoutSecs: 5
@@ -131,7 +181,20 @@ describe('run-queries', () => {
       })),
       logger: {
         log: sandbox.spy()
-      }
-    };
+      },
+      cliServer
+    } as unknown as QueryServerClient;
+  }
+
+  function createMockCliServer(mockOperations: Record<string, any[]>): CodeQLCliServer {
+    const mockServer: Record<string, any> = {};
+    for (const [operation, returns] of Object.entries(mockOperations)) {
+      mockServer[operation] = sandbox.stub();
+      returns.forEach((returnValue, i) => {
+        mockServer[operation].onCall(i).resolves(returnValue);
+      });
+    }
+
+    return mockServer as unknown as CodeQLCliServer;
   }
 });