Add support for automodel requests in the mock API server (#2674)

Author: charisk

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/gh-api-request.ts

@@ -11,6 +11,7 @@ export enum RequestKind {
   GetVariantAnalysisRepo = "getVariantAnalysisRepo",
   GetVariantAnalysisRepoResult = "getVariantAnalysisRepoResult",
   CodeSearch = "codeSearch",
+  AutoModel = "autoModel",
 }
 
 interface BasicErorResponse {
@@ -87,13 +88,27 @@ interface CodeSearchRequest {
   };
 }
 
+interface AutoModelRequest {
+  request: {
+    kind: RequestKind.AutoModel;
+  };
+  response: {
+    status: number;
+    body?: {
+      candidates: string;
+    };
+    message?: string;
+  };
+}
+
 export type GitHubApiRequest =
   | GetRepoRequest
   | SubmitVariantAnalysisRequest
   | GetVariantAnalysisRequest
   | GetVariantAnalysisRepoRequest
   | GetVariantAnalysisRepoResultRequest
-  | CodeSearchRequest;
+  | CodeSearchRequest
+  | AutoModelRequest;
 
 export const isGetRepoRequest = (
   request: GitHubApiRequest,
@@ -123,3 +138,8 @@ export const isCodeSearchRequest = (
   request: GitHubApiRequest,
 ): request is CodeSearchRequest =>
   request.request.kind === RequestKind.CodeSearch;
+
+export const isAutoModelRequest = (
+  request: GitHubApiRequest,
+): request is AutoModelRequest =>
+  request.request.kind === RequestKind.AutoModel;

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/recorder.ts

@@ -259,6 +259,21 @@ async function createGitHubApiRequest(
     };
   }
 
+  const autoModelMatch = url.match(
+    /\/repos\/github\/codeql\/code-scanning\/codeql\/auto-model/,
+  );
+  if (autoModelMatch) {
+    return {
+      request: {
+        kind: RequestKind.AutoModel,
+      },
+      response: {
+        status,
+        body: JSON.parse(body),
+      },
+    };
+  }
+
   return undefined;
 }
 

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/request-handlers.ts

@@ -3,6 +3,7 @@ import { readdir, readJson, readFile } from "fs-extra";
 import { DefaultBodyType, MockedRequest, rest, RestHandler } from "msw";
 import {
   GitHubApiRequest,
+  isAutoModelRequest,
   isCodeSearchRequest,
   isGetRepoRequest,
   isGetVariantAnalysisRepoRequest,
@@ -27,6 +28,7 @@ export async function createRequestHandlers(
     createGetVariantAnalysisRepoRequestHandler(requests),
     createGetVariantAnalysisRepoResultRequestHandler(requests),
     createCodeSearchRequestHandler(requests),
+    createAutoModelRequestHandler(requests),
   ];
 
   return handlers;
@@ -219,3 +221,30 @@ function createCodeSearchRequestHandler(
     );
   });
 }
+
+function createAutoModelRequestHandler(
+  requests: GitHubApiRequest[],
+): RequestHandler {
+  const autoModelRequests = requests.filter(isAutoModelRequest);
+  let requestIndex = 0;
+
+  // During automodeling there can be multiple API requests for each batch
+  // of candidates we want to model. We need to return different responses for each request,
+  // so keep an index of the request and return the appropriate response.
+  return rest.get(
+    `${baseUrl}/code-scanning/codeql/auto-model`,
+    (_req, res, ctx) => {
+      const request = autoModelRequests[requestIndex];
+
+      if (requestIndex < autoModelRequests.length - 1) {
+        // If there are more requests to come, increment the index.
+        requestIndex++;
+      }
+
+      return res(
+        ctx.status(request.response.status),
+        ctx.json(request.response.body),
+      );
+    },
+  );
+}