From 7e5953dec1b807e0d6873c2fcc25ea4d64bb416b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Vannicatte?= <vannicattec@gmail.com>
Date: Thu, 19 May 2022 15:37:53 +0200
Subject: [PATCH] fix(scripts): codegen cleanup

---
 .github/workflows/codegen.yml                 |  1 +
 scripts/ci/codegen/__tests__/codegen.test.ts  | 49 ++++++++++++++-----
 scripts/ci/codegen/text.ts                    |  5 +-
 scripts/ci/codegen/upsertGenerationComment.ts | 10 +++-
 4 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/codegen.yml b/.github/workflows/codegen.yml
index 0cf464513b..787da8c56a 100644
--- a/.github/workflows/codegen.yml
+++ b/.github/workflows/codegen.yml
@@ -25,6 +25,7 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.TOKEN_GENERATE_BOT }}
           PR_NUMBER: ${{ github.event.number }}
+          HEAD_BRANCH: ${{ github.head_ref }}
 
   cleanup:
     runs-on: ubuntu-20.04
diff --git a/scripts/ci/codegen/__tests__/codegen.test.ts b/scripts/ci/codegen/__tests__/codegen.test.ts
index 5ab99b270c..6d482f4e7b 100644
--- a/scripts/ci/codegen/__tests__/codegen.test.ts
+++ b/scripts/ci/codegen/__tests__/codegen.test.ts
@@ -1,17 +1,12 @@
-import { MAIN_BRANCH } from '../../../common';
+import * as common from '../../../common';
 import { cleanGeneratedBranch } from '../cleanGeneratedBranch';
 import { pushGeneratedCode } from '../pushGeneratedCode';
 import commentText from '../text';
 import {
-  getCommentBody,
   upsertGenerationComment,
+  getCommentBody,
 } from '../upsertGenerationComment';
 
-jest.mock('../../../common', () => ({
-  ...(jest.requireActual('../../../common') as any),
-  run: jest.fn().mockResolvedValue('mocked'),
-}));
-
 describe('codegen', () => {
   describe('cleanGeneratedBranch', () => {
     it('throws without parameters', async () => {
@@ -68,13 +63,41 @@ describe('codegen', () => {
       `);
     });
 
-    it('returns the right comment for a `cleanup` trigger', async () => {
-      expect(await getCommentBody('cleanup')).toMatchInlineSnapshot(`
-        "### ✗ The generated branch has been deleted.
+    describe('cleanup', () => {
+      let mockedResolvedValue: string;
+      beforeEach(() => {
+        jest.spyOn(common, 'run').mockImplementation(() => {
+          return Promise.resolve(mockedResolvedValue);
+        });
+      });
+
+      afterEach(() => {
+        jest.spyOn(common, 'run').mockRestore();
+      });
+
+      afterEach(() => {});
+      it('returns the right comment for a `cleanup` trigger', async () => {
+        mockedResolvedValue = 'mocked';
 
-        If the PR has been merged, you can check the generated code on the [\`${MAIN_BRANCH}\` branch](https://github.com/algolia/api-clients-automation/tree/${MAIN_BRANCH}).
-        You can still access [the last generated commit](https://github.com/algolia/api-clients-automation/commit/mocked)."
-      `);
+        expect(await getCommentBody('cleanup')).toMatchInlineSnapshot(`
+          "### ✗ The generated branch has been deleted.
+
+          If the PR has been merged, you can check the generated code on the [\`${common.MAIN_BRANCH}\` branch](https://github.com/algolia/api-clients-automation/tree/${common.MAIN_BRANCH}).
+          You can still access the code generated on \`mocked\` via [this commit](https://github.com/algolia/api-clients-automation/commit/mocked)."
+          `);
+      });
+
+      it('fallbacks to the env variable HEAD_BRANCH if found when we are on `main`', async () => {
+        process.env.HEAD_BRANCH = 'myFakeBranch';
+        mockedResolvedValue = 'main';
+
+        expect(await getCommentBody('cleanup')).toMatchInlineSnapshot(`
+          "### ✗ The generated branch has been deleted.
+
+          If the PR has been merged, you can check the generated code on the [\`${common.MAIN_BRANCH}\` branch](https://github.com/algolia/api-clients-automation/tree/${common.MAIN_BRANCH}).
+          You can still access the code generated on \`generated/myFakeBranch\` via [this commit](https://github.com/algolia/api-clients-automation/commit/main)."
+        `);
+      });
     });
 
     describe('text', () => {
diff --git a/scripts/ci/codegen/text.ts b/scripts/ci/codegen/text.ts
index 6ac416e138..7de7ddb749 100644
--- a/scripts/ci/codegen/text.ts
+++ b/scripts/ci/codegen/text.ts
@@ -15,9 +15,10 @@ export default {
   cleanup: {
     header: '### ✗ The generated branch has been deleted.',
     body: (
-      generatedCommit: string
+      generatedCommit: string,
+      branch: string
     ): string => `If the PR has been merged, you can check the generated code on the [\`${MAIN_BRANCH}\` branch](${REPO_URL}/tree/${MAIN_BRANCH}).
-You can still access [the last generated commit](${REPO_URL}/commit/${generatedCommit}).`,
+You can still access the code generated on \`${branch}\` via [this commit](${REPO_URL}/commit/${generatedCommit}).`,
   },
   codegen: {
     header: '### ✔️ Code generated!',
diff --git a/scripts/ci/codegen/upsertGenerationComment.ts b/scripts/ci/codegen/upsertGenerationComment.ts
index f2d6c84a7a..b9457566b7 100644
--- a/scripts/ci/codegen/upsertGenerationComment.ts
+++ b/scripts/ci/codegen/upsertGenerationComment.ts
@@ -19,10 +19,16 @@ const allowedTriggers = [
 type Trigger = typeof allowedTriggers[number];
 
 export async function getCommentBody(trigger: Trigger): Promise<string> {
-  const generatedBranch = await run('git branch --show-current');
+  let generatedBranch = await run('git branch --show-current');
+
+  // `cleanup` is triggered on PR close, which runs on `main`, so we lose the
+  // branch name context at this point
+  if (generatedBranch === 'main' && process.env.HEAD_BRANCH) {
+    generatedBranch = `generated/${process.env.HEAD_BRANCH}`;
+  }
+
   const baseBranch = generatedBranch.replace('generated/', '');
   const baseCommit = await run(`git show ${baseBranch} -s --format=%H`);
-
   const generatedCommit = await run(
     `git show ${generatedBranch} -s --format=%H`
   );