fix(github-upload-public-key): resolve issues with flaky tests (#634)

## Description

Better test cleanup, and resolve flakiness.
<!-- Briefly describe what this PR does and why -->

## Type of Change

- [ ] New module
- [ ] New template
- [X] Bug fix
- [ ] Feature/enhancement
- [ ] Documentation
- [ ] Other

## Module Information

<!-- Delete this section if not applicable -->

**Path:** `registry/coder/modules/github-upload-public-key`  
**New version:** `v1.0.32`  
**Breaking change:** [ ] Yes [ ] No

## Testing & Validation

- [X] Tests pass (`bun test`)
- [X] Code formatted (`bun fmt`)
- [X] Changes tested locally

## Related Issues

<!-- Link related issues or write "None" if not applicable -->

registry/coder/modules/github-upload-public-key/main.test.ts

@@ -1,9 +1,17 @@
-import { type Server, serve } from "bun";
-import { describe, expect, it } from "bun:test";
+import { serve } from "bun";
+import {
+  afterEach,
+  beforeAll,
+  describe,
+  expect,
+  it,
+  setDefaultTimeout,
+} from "bun:test";
 import {
   createJSONResponse,
   execContainer,
   findResourceInstance,
+  removeContainer,
   runContainer,
   runTerraformApply,
   runTerraformInit,
@@ -11,77 +19,48 @@ import {
   writeCoder,
 } from "~test";
 
-describe("github-upload-public-key", async () => {
-  await runTerraformInit(import.meta.dir);
-
-  testRequiredVariables(import.meta.dir, {
-    agent_id: "foo",
-  });
-
-  it("creates new key if one does not exist", async () => {
-    const { instance, id, server } = await setupContainer();
-    await writeCoder(id, "echo foo");
-
-    const url = server.url.toString().slice(0, -1);
-    const exec = await execContainer(id, [
-      "env",
-      `CODER_ACCESS_URL=${url}`,
-      `GITHUB_API_URL=${url}`,
-      "CODER_OWNER_SESSION_TOKEN=foo",
-      "CODER_EXTERNAL_AUTH_ID=github",
-      "bash",
-      "-c",
-      instance.script,
-    ]);
-    expect(exec.stdout).toContain(
-      "Your Coder public key has been added to GitHub!",
-    );
-    expect(exec.exitCode).toBe(0);
-    // we need to increase timeout to pull the container
-  }, 15000);
-
-  it("does nothing if one already exists", async () => {
-    const { instance, id, server } = await setupContainer();
-    // use keyword to make server return a existing key
-    await writeCoder(id, "echo findkey");
-
-    const url = server.url.toString().slice(0, -1);
-    const exec = await execContainer(id, [
-      "env",
-      `CODER_ACCESS_URL=${url}`,
-      `GITHUB_API_URL=${url}`,
-      "CODER_OWNER_SESSION_TOKEN=foo",
-      "CODER_EXTERNAL_AUTH_ID=github",
-      "bash",
-      "-c",
-      instance.script,
-    ]);
-    expect(exec.stdout).toContain(
-      "Your Coder public key is already on GitHub!",
-    );
-    expect(exec.exitCode).toBe(0);
-  });
+let cleanupFunctions: (() => Promise<void>)[] = [];
+const registerCleanup = (cleanup: () => Promise<void>) => {
+  cleanupFunctions.push(cleanup);
+};
+afterEach(async () => {
+  const cleanupFnsCopy = cleanupFunctions.slice().reverse();
+  cleanupFunctions = [];
+  for (const cleanup of cleanupFnsCopy) {
+    try {
+      await cleanup();
+    } catch (error) {
+      console.error("Error during cleanup:", error);
+    }
+  }
 });
 
 const setupContainer = async (
   image = "lorello/alpine-bash",
   vars: Record<string, string> = {},
 ) => {
-  const server = await setupServer();
+  const server = setupServer();
   const state = await runTerraformApply(import.meta.dir, {
     agent_id: "foo",
     ...vars,
   });
   const instance = findResourceInstance(state, "coder_script");
   const id = await runContainer(image);
+
+  registerCleanup(async () => {
+    server.stop();
+  });
+  registerCleanup(async () => {
+    await removeContainer(id);
+  });
+
   return { id, instance, server };
 };
 
-const setupServer = async (): Promise<Server> => {
-  let url: URL;
-  const fakeSlackHost = serve({
+const setupServer = () => {
+  const fakeGithubHost = serve({
     fetch: (req) => {
-      url = new URL(req.url);
+      const url = new URL(req.url);
       if (url.pathname === "/api/v2/users/me/gitsshkey") {
         return createJSONResponse({
           public_key: "exists",
@@ -128,5 +107,60 @@ const setupServer = async (): Promise<Server> => {
     port: 0,
   });
 
-  return fakeSlackHost;
+  return fakeGithubHost;
 };
+
+setDefaultTimeout(30 * 1000);
+
+describe("github-upload-public-key", () => {
+  beforeAll(async () => {
+    await runTerraformInit(import.meta.dir);
+  });
+
+  testRequiredVariables(import.meta.dir, {
+    agent_id: "foo",
+  });
+
+  it("creates new key if one does not exist", async () => {
+    const { instance, id, server } = await setupContainer();
+    await writeCoder(id, "echo foo");
+
+    const url = server.url.toString().slice(0, -1);
+    const exec = await execContainer(id, [
+      "env",
+      `CODER_ACCESS_URL=${url}`,
+      `GITHUB_API_URL=${url}`,
+      "CODER_OWNER_SESSION_TOKEN=foo",
+      "CODER_EXTERNAL_AUTH_ID=github",
+      "bash",
+      "-c",
+      instance.script,
+    ]);
+    expect(exec.stdout).toContain(
+      "Your Coder public key has been added to GitHub!",
+    );
+    expect(exec.exitCode).toBe(0);
+  });
+
+  it("does nothing if one already exists", async () => {
+    const { instance, id, server } = await setupContainer();
+    // use keyword to make server return a existing key
+    await writeCoder(id, "echo findkey");
+
+    const url = server.url.toString().slice(0, -1);
+    const exec = await execContainer(id, [
+      "env",
+      `CODER_ACCESS_URL=${url}`,
+      `GITHUB_API_URL=${url}`,
+      "CODER_OWNER_SESSION_TOKEN=foo",
+      "CODER_EXTERNAL_AUTH_ID=github",
+      "bash",
+      "-c",
+      instance.script,
+    ]);
+    expect(exec.stdout).toContain(
+      "Your Coder public key is already on GitHub!",
+    );
+    expect(exec.exitCode).toBe(0);
+  });
+});