Add tests for clear cleanup size telemetry

Author: henrymercer

src/database-upload.test.ts

@@ -11,8 +11,10 @@ import * as apiClient from "./api-client";
 import { createStubCodeQL } from "./codeql";
 import { Config } from "./config-utils";
 import { cleanupAndUploadDatabases } from "./database-upload";
+import { Feature } from "./feature-flags";
 import * as gitUtils from "./git-utils";
 import { BuiltInLanguage } from "./languages";
+import { OverlayDatabaseMode } from "./overlay/overlay-database-mode";
 import { RepositoryNwo } from "./repository";
 import {
   checkExpectedLogMessages,
@@ -24,6 +26,7 @@ import {
   setupTests,
 } from "./testing-utils";
 import {
+  CleanupLevel,
   GitHubVariant,
   HTTPError,
   initializeEnvironment,
@@ -335,3 +338,146 @@ test.serial("Successfully uploading a database to GHEC-DR", async (t) => {
     );
   });
 });
+
+test.serial(
+  "Records overlay and clear cleanup sizes when uploading an overlay-base database",
+  async (t) => {
+    await withTmpDir(async (tmpDir) => {
+      setupActionsVars(tmpDir, tmpDir);
+      sinon
+        .stub(actionsUtil, "getRequiredInput")
+        .withArgs("upload-database")
+        .returns("true");
+      sinon.stub(gitUtils, "isAnalyzingDefaultBranch").resolves(true);
+
+      await mockHttpRequests(201);
+
+      // Track the cleanup level passed to each cleanup so that the database
+      // bundle stub can write a differently-sized bundle for each level.
+      const cleanupLevels: CleanupLevel[] = [];
+      let lastCleanupLevel: CleanupLevel | undefined;
+      const overlaySizeBytes = 100;
+      const clearSizeBytes = 50;
+      const codeql = createStubCodeQL({
+        async databaseCleanupCluster(_config, cleanupLevel) {
+          cleanupLevels.push(cleanupLevel);
+          lastCleanupLevel = cleanupLevel;
+        },
+        async databaseBundle(_databasePath, outputFilePath) {
+          const sizeBytes =
+            lastCleanupLevel === CleanupLevel.Overlay
+              ? overlaySizeBytes
+              : clearSizeBytes;
+          fs.writeFileSync(outputFilePath, "x".repeat(sizeBytes));
+        },
+      });
+
+      const config = getTestConfig(tmpDir);
+      config.overlayDatabaseMode = OverlayDatabaseMode.OverlayBase;
+
+      const loggedMessages: LoggedMessage[] = [];
+      const results = await cleanupAndUploadDatabases(
+        testRepoName,
+        codeql,
+        config,
+        testApiDetails,
+        createFeatures([Feature.UploadOverlayDbToApi]),
+        getRecordingLogger(loggedMessages),
+      );
+
+      // The database should be cleaned up at the `overlay` level for the upload
+      // and then re-cleaned at the `clear` level to measure its size.
+      t.deepEqual(cleanupLevels, [CleanupLevel.Overlay, CleanupLevel.Clear]);
+
+      t.is(results.length, 1);
+      t.is(results[0].is_overlay_base, true);
+      t.is(results[0].zipped_upload_size_bytes, overlaySizeBytes);
+      t.is(results[0].clear_cleanup_zipped_size_bytes, clearSizeBytes);
+      t.is(typeof results[0].clear_cleanup_measurement_duration_ms, "number");
+    });
+  },
+);
+
+test.serial(
+  "Does not measure clear cleanup size for a regular (non-overlay-base) upload",
+  async (t) => {
+    await withTmpDir(async (tmpDir) => {
+      setupActionsVars(tmpDir, tmpDir);
+      sinon
+        .stub(actionsUtil, "getRequiredInput")
+        .withArgs("upload-database")
+        .returns("true");
+      sinon.stub(gitUtils, "isAnalyzingDefaultBranch").resolves(true);
+
+      await mockHttpRequests(201);
+
+      const cleanupLevels: CleanupLevel[] = [];
+      const codeql = createStubCodeQL({
+        async databaseCleanupCluster(_config, cleanupLevel) {
+          cleanupLevels.push(cleanupLevel);
+        },
+        async databaseBundle(_databasePath, outputFilePath) {
+          fs.writeFileSync(outputFilePath, "");
+        },
+      });
+
+      const results = await cleanupAndUploadDatabases(
+        testRepoName,
+        codeql,
+        getTestConfig(tmpDir),
+        testApiDetails,
+        createFeatures([Feature.UploadOverlayDbToApi]),
+        getRecordingLogger([]),
+      );
+
+      // A regular upload is cleaned only once, at the `clear` level.
+      t.deepEqual(cleanupLevels, [CleanupLevel.Clear]);
+      t.is(results[0].is_overlay_base, false);
+      t.is(results[0].clear_cleanup_zipped_size_bytes, undefined);
+      t.is(results[0].clear_cleanup_measurement_duration_ms, undefined);
+    });
+  },
+);
+
+test.serial("Does not measure clear cleanup size in debug mode", async (t) => {
+  await withTmpDir(async (tmpDir) => {
+    setupActionsVars(tmpDir, tmpDir);
+    sinon
+      .stub(actionsUtil, "getRequiredInput")
+      .withArgs("upload-database")
+      .returns("true");
+    sinon.stub(gitUtils, "isAnalyzingDefaultBranch").resolves(true);
+
+    await mockHttpRequests(201);
+
+    const cleanupLevels: CleanupLevel[] = [];
+    const codeql = createStubCodeQL({
+      async databaseCleanupCluster(_config, cleanupLevel) {
+        cleanupLevels.push(cleanupLevel);
+      },
+      async databaseBundle(_databasePath, outputFilePath) {
+        fs.writeFileSync(outputFilePath, "");
+      },
+    });
+
+    const config = getTestConfig(tmpDir);
+    config.overlayDatabaseMode = OverlayDatabaseMode.OverlayBase;
+    config.debugMode = true;
+
+    const results = await cleanupAndUploadDatabases(
+      testRepoName,
+      codeql,
+      config,
+      testApiDetails,
+      createFeatures([Feature.UploadOverlayDbToApi]),
+      getRecordingLogger([]),
+    );
+
+    // In debug mode we clean up at the `overlay` level for the upload but skip
+    // the additional `clear` cleanup, to preserve the database for debugging.
+    t.deepEqual(cleanupLevels, [CleanupLevel.Overlay]);
+    t.is(results[0].is_overlay_base, true);
+    t.is(results[0].clear_cleanup_zipped_size_bytes, undefined);
+    t.is(results[0].clear_cleanup_measurement_duration_ms, undefined);
+  });
+});