From 8c5122ea75fe7e9f2a8883731dcd14144ea2512e Mon Sep 17 00:00:00 2001
From: Chuan-kai Lin <cklin@github.com>
Date: Mon, 7 Jul 2025 08:13:06 -0700
Subject: [PATCH] Add getPullRequestBranches() tests

---
 src/actions-util.test.ts | 131 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 131 insertions(+)

diff --git a/src/actions-util.test.ts b/src/actions-util.test.ts
index f4809d759..f5be61a0d 100644
--- a/src/actions-util.test.ts
+++ b/src/actions-util.test.ts
@@ -1,5 +1,7 @@
+import * as github from "@actions/github";
 import test from "ava";
 
+import { getPullRequestBranches, isAnalyzingPullRequest } from "./actions-util";
 import { computeAutomationID } from "./api-client";
 import { EnvVar } from "./environment";
 import { setupTests } from "./testing-utils";
@@ -7,6 +9,39 @@ import { initializeEnvironment } from "./util";
 
 setupTests(test);
 
+function withMockedContext<T>(mockPayload: any, testFn: () => T): T {
+  const originalPayload = github.context.payload;
+  github.context.payload = mockPayload;
+  try {
+    return testFn();
+  } finally {
+    github.context.payload = originalPayload;
+  }
+}
+
+function withMockedEnv<T>(
+  envVars: Record<string, string | undefined>,
+  testFn: () => T,
+): T {
+  const originalEnv = { ...process.env };
+
+  // Apply environment changes
+  for (const [key, value] of Object.entries(envVars)) {
+    if (value === undefined) {
+      delete process.env[key];
+    } else {
+      process.env[key] = value;
+    }
+  }
+
+  try {
+    return testFn();
+  } finally {
+    // Restore original environment
+    process.env = originalEnv;
+  }
+}
+
 test("computeAutomationID()", async (t) => {
   let actualAutomationID = computeAutomationID(
     ".github/workflows/codeql-analysis.yml:analyze",
@@ -58,6 +93,102 @@ test("computeAutomationID()", async (t) => {
   );
 });
 
+test("getPullRequestBranches() with pull request context", (t) => {
+  withMockedContext(
+    {
+      pull_request: {
+        number: 123,
+        base: { ref: "main" },
+        head: { label: "user:feature-branch" },
+      },
+    },
+    () => {
+      t.deepEqual(getPullRequestBranches(), {
+        base: "main",
+        head: "user:feature-branch",
+      });
+      t.is(isAnalyzingPullRequest(), true);
+    },
+  );
+});
+
+test("getPullRequestBranches() returns undefined with push context", (t) => {
+  withMockedContext(
+    {
+      push: {
+        ref: "refs/heads/main",
+      },
+    },
+    () => {
+      t.is(getPullRequestBranches(), undefined);
+      t.is(isAnalyzingPullRequest(), false);
+    },
+  );
+});
+
+test("getPullRequestBranches() with Default Setup environment variables", (t) => {
+  withMockedContext({}, () => {
+    withMockedEnv(
+      {
+        CODE_SCANNING_REF: "refs/heads/feature-branch",
+        CODE_SCANNING_BASE_BRANCH: "main",
+      },
+      () => {
+        t.deepEqual(getPullRequestBranches(), {
+          base: "main",
+          head: "refs/heads/feature-branch",
+        });
+        t.is(isAnalyzingPullRequest(), true);
+      },
+    );
+  });
+});
+
+test("getPullRequestBranches() returns undefined when only CODE_SCANNING_REF is set", (t) => {
+  withMockedContext({}, () => {
+    withMockedEnv(
+      {
+        CODE_SCANNING_REF: "refs/heads/feature-branch",
+        CODE_SCANNING_BASE_BRANCH: undefined,
+      },
+      () => {
+        t.is(getPullRequestBranches(), undefined);
+        t.is(isAnalyzingPullRequest(), false);
+      },
+    );
+  });
+});
+
+test("getPullRequestBranches() returns undefined when only CODE_SCANNING_BASE_BRANCH is set", (t) => {
+  withMockedContext({}, () => {
+    withMockedEnv(
+      {
+        CODE_SCANNING_REF: undefined,
+        CODE_SCANNING_BASE_BRANCH: "main",
+      },
+      () => {
+        t.is(getPullRequestBranches(), undefined);
+        t.is(isAnalyzingPullRequest(), false);
+      },
+    );
+  });
+});
+
+test("getPullRequestBranches() returns undefined when no PR context", (t) => {
+  withMockedContext({}, () => {
+    withMockedEnv(
+      {
+        CODE_SCANNING_REF: undefined,
+        CODE_SCANNING_BASE_BRANCH: undefined,
+      },
+      () => {
+        t.is(getPullRequestBranches(), undefined);
+        t.is(isAnalyzingPullRequest(), false);
+      },
+    );
+  });
+});
+
 test("initializeEnvironment", (t) => {
   initializeEnvironment("1.2.3");
   t.deepEqual(process.env[EnvVar.VERSION], "1.2.3");