Add tests for repo states

extensions/ql-vscode/src/remote-queries/variant-analysis-manager.ts

@@ -79,7 +79,7 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
       if (status === QueryStatus.InProgress) {
         // In this case, last time we checked, the query was still in progress.
         // We need to setup the monitor to check for completion.
-        await commands.executeCommand('codeQL.monitorVariantAnalysis', variantAnalysis);
+        void commands.executeCommand('codeQL.monitorVariantAnalysis', variantAnalysis);
       }
     }
   }
@@ -241,7 +241,13 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
       repoState.downloadStatus = VariantAnalysisScannedRepositoryDownloadStatus.InProgress;
       await this.onRepoStateUpdated(variantAnalysisSummary.id, repoState);
 
-      await this.variantAnalysisResultsManager.download(credentials, variantAnalysisSummary.id, repoTask, this.getVariantAnalysisStorageLocation(variantAnalysisSummary.id));
+      try {
+        await this.variantAnalysisResultsManager.download(credentials, variantAnalysisSummary.id, repoTask, this.getVariantAnalysisStorageLocation(variantAnalysisSummary.id));
+      } catch (e) {
+        repoState.downloadStatus = VariantAnalysisScannedRepositoryDownloadStatus.Failed;
+        await this.onRepoStateUpdated(variantAnalysisSummary.id, repoState);
+        throw new Error(`Could not download the results for variant analysis with id: ${variantAnalysisSummary.id}. Error: ${getErrorMessage(e)}`);
+      }
     }
 
     repoState.downloadStatus = VariantAnalysisScannedRepositoryDownloadStatus.Succeeded;

extensions/ql-vscode/src/vscode-tests/cli-integration/remote-queries/variant-analysis-manager.test.ts

@@ -21,11 +21,18 @@ import { createMockVariantAnalysisRepoTask } from '../../factories/remote-querie
 import { CodeQLCliServer } from '../../../cli';
 import { storagePath } from '../global.helper';
 import { VariantAnalysisResultsManager } from '../../../remote-queries/variant-analysis-results-manager';
-import { VariantAnalysis } from '../../../remote-queries/shared/variant-analysis';
+import {
+  VariantAnalysis,
+  VariantAnalysisScannedRepositoryDownloadStatus
+} from '../../../remote-queries/shared/variant-analysis';
 import { createMockVariantAnalysis } from '../../factories/remote-queries/shared/variant-analysis';
+import { QueryStatus } from '../../../query-status';
 
-describe('Variant Analysis Manager', async function() {
+describe.only('Variant Analysis Manager', async function() {
   let sandbox: sinon.SinonSandbox;
+  let pathExistsStub: sinon.SinonStub;
+  let readJsonStub: sinon.SinonStub;
+  let outputJsonStub: sinon.SinonStub;
   let cli: CodeQLCliServer;
   let cancellationTokenSource: CancellationTokenSource;
   let variantAnalysisManager: VariantAnalysisManager;
@@ -41,6 +48,9 @@ describe('Variant Analysis Manager', async function() {
     sandbox.stub(config, 'isVariantAnalysisLiveResultsEnabled').returns(false);
     sandbox.stub(fs, 'mkdirSync');
     sandbox.stub(fs, 'writeFile');
+    pathExistsStub = sandbox.stub(fs, 'pathExists').callThrough();
+    readJsonStub = sandbox.stub(fs, 'readJson').callThrough();
+    outputJsonStub = sandbox.stub(fs, 'outputJson');
 
     cancellationTokenSource = new CancellationTokenSource();
 
@@ -61,6 +71,74 @@ describe('Variant Analysis Manager', async function() {
     sandbox.restore();
   });
 
+  describe('rehydrateVariantAnalysis', () => {
+    const variantAnalysis = createMockVariantAnalysis();
+
+    describe('when the directory does not exist', () => {
+      beforeEach(() => {
+        pathExistsStub.withArgs(path.join(storagePath, variantAnalysis.id.toString())).resolves(false);
+      });
+
+      it('should fire the removed event if the file does not exist', async () => {
+        const stub = sandbox.stub();
+        variantAnalysisManager.onVariantAnalysisRemoved(stub);
+
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis, QueryStatus.InProgress);
+
+        expect(stub).to.have.been.calledOnce;
+        sinon.assert.calledWith(pathExistsStub, path.join(storagePath, variantAnalysis.id.toString()));
+      });
+    });
+
+    describe('when the directory exists', () => {
+      beforeEach(() => {
+        pathExistsStub.withArgs(path.join(storagePath, variantAnalysis.id.toString())).resolves(true);
+      });
+
+      it('should store the variant analysis', async () => {
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis, QueryStatus.InProgress);
+
+        expect(await variantAnalysisManager.getVariantAnalysis(variantAnalysis.id)).to.deep.equal(variantAnalysis);
+      });
+
+      it('should not error if the repo states file does not exist', async () => {
+        readJsonStub.withArgs(path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json')).rejects(new Error('File does not exist'));
+
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis, QueryStatus.InProgress);
+
+        sinon.assert.calledWith(readJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'));
+        expect(await variantAnalysisManager.getRepoStates(variantAnalysis.id)).to.deep.equal([]);
+      });
+
+      it('should read in the repo states if it exists', async () => {
+        readJsonStub.withArgs(path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json')).resolves({
+          [scannedRepos[0].repository.id]: {
+            repositoryId: scannedRepos[0].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+          },
+          [scannedRepos[1].repository.id]: {
+            repositoryId: scannedRepos[1].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+          },
+        });
+
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis, QueryStatus.InProgress);
+
+        sinon.assert.calledWith(readJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'));
+        expect(await variantAnalysisManager.getRepoStates(variantAnalysis.id)).to.have.same.deep.members([
+          {
+            repositoryId: scannedRepos[0].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+          },
+          {
+            repositoryId: scannedRepos[1].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+          },
+        ]);
+      });
+    });
+  });
+
   describe('when credentials are invalid', async () => {
     beforeEach(async () => { sandbox.stub(Credentials, 'initialize').resolves(undefined); });
 
@@ -155,6 +233,171 @@ describe('Variant Analysis Manager', async function() {
 
           expect(getVariantAnalysisRepoResultStub.calledOnce).to.be.true;
         });
+
+        it('should skip the download if the repository has already been downloaded', async () => {
+          // First, do a download so it is downloaded. This avoids having to mock the repo states.
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          getVariantAnalysisRepoStub.resetHistory();
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          expect(getVariantAnalysisRepoStub.notCalled).to.be.true;
+        });
+
+        it('should write the repo state when the download is successful', async () => {
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+          });
+        });
+
+        it('should not write the repo state when the download fails', async () => {
+          getVariantAnalysisRepoResultStub.rejects(new Error('Failed to download'));
+
+          try {
+            await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+              scannedRepos[0],
+              variantAnalysis,
+              cancellationTokenSource.token
+            );
+            fail('Expected an error to be thrown');
+          } catch (e: any) {
+            // we can ignore this error, we expect this
+          }
+
+          sinon.assert.notCalled(outputJsonStub);
+        });
+
+        it('should have a failed repo state when the repo task API fails', async () => {
+          getVariantAnalysisRepoStub.onFirstCall().rejects(new Error('Failed to download'));
+
+          try {
+            await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+              scannedRepos[0],
+              variantAnalysis,
+              cancellationTokenSource.token
+            );
+            fail('Expected an error to be thrown');
+          } catch (e) {
+            // we can ignore this error, we expect this
+          }
+
+          sinon.assert.notCalled(outputJsonStub);
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[1],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Failed,
+            },
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+          });
+        });
+
+        it('should have a failed repo state when the download fails', async () => {
+          getVariantAnalysisRepoResultStub.onFirstCall().rejects(new Error('Failed to download'));
+
+          try {
+            await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+              scannedRepos[0],
+              variantAnalysis,
+              cancellationTokenSource.token
+            );
+            fail('Expected an error to be thrown');
+          } catch (e) {
+            // we can ignore this error, we expect this
+          }
+
+          sinon.assert.notCalled(outputJsonStub);
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[1],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Failed,
+            },
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+          });
+        });
+
+        it('should update the repo state correctly', async () => {
+          // To set some initial repo states, we need to mock the correct methods so that the repo states are read in.
+          // The actual tests for these are in rehydrateVariantAnalysis, so we can just mock them here and test that
+          // the methods are called.
+
+          pathExistsStub.withArgs(path.join(storagePath, variantAnalysis.id.toString())).resolves(true);
+          // This will read in the correct repo states
+          readJsonStub.withArgs(path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json')).resolves({
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+            [scannedRepos[2].repository.id]: {
+              repositoryId: scannedRepos[2].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+            },
+          });
+
+          await variantAnalysisManager.rehydrateVariantAnalysis({
+            ...createMockVariantAnalysis(),
+            id: variantAnalysis.id,
+          }, QueryStatus.InProgress);
+          sinon.assert.calledWith(readJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'));
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+            [scannedRepos[2].repository.id]: {
+              repositoryId: scannedRepos[2].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+            },
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            }
+          });
+        });
       });
 
       describe('enqueueDownload', async () => {