diff --git a/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.test.ts b/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.test.ts
new file mode 100644
index 0000000000000..ec5d24463a3a6
--- /dev/null
+++ b/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.test.ts
@@ -0,0 +1,37 @@
+import { ProvisionerJobLog } from "api/typesGenerated"
+import { groupLogsByStage } from "./WorkspaceBuildLogs"
+
+describe("groupLogsByStage", () => {
+  it("should group them by stage", () => {
+    const input: ProvisionerJobLog[] = [
+      {
+        id: "1",
+        created_at: "oct 13",
+        log_source: "provisioner",
+        log_level: "debug",
+        stage: "build",
+        output: "test",
+      },
+      {
+        id: "2",
+        created_at: "oct 13",
+        log_source: "provisioner",
+        log_level: "debug",
+        stage: "cleanup",
+        output: "test",
+      },
+      {
+        id: "3",
+        created_at: "oct 13",
+        log_source: "provisioner",
+        log_level: "debug",
+        stage: "cleanup",
+        output: "done",
+      },
+    ]
+
+    const actual = groupLogsByStage(input)
+
+    expect(actual["cleanup"].length).toBe(2)
+  })
+})
diff --git a/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.tsx b/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.tsx
index cc1b9e0233601..f0dc70c5609bc 100644
--- a/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.tsx
+++ b/site/src/components/WorkspaceBuildLogs/WorkspaceBuildLogs.tsx
@@ -10,18 +10,18 @@ const Language = {
 }
 
 type Stage = ProvisionerJobLog["stage"]
+type LogsGroupedByStage = Record<Stage, ProvisionerJobLog[]>
+type GroupLogsByStageFn = (logs: ProvisionerJobLog[]) => LogsGroupedByStage
 
-const groupLogsByStage = (logs: ProvisionerJobLog[]) => {
-  const logsByStage: Record<Stage, ProvisionerJobLog[]> = {}
+export const groupLogsByStage: GroupLogsByStageFn = (logs) => {
+  const logsByStage: LogsGroupedByStage = {}
 
   for (const log of logs) {
-    // If there is no log in the stage record, add an empty array
-    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-    if (logsByStage[log.stage] === undefined) {
-      logsByStage[log.stage] = []
+    if (log.stage in logsByStage) {
+      logsByStage[log.stage].push(log)
+    } else {
+      logsByStage[log.stage] = [log]
     }
-
-    logsByStage[log.stage].push(log)
   }
 
   return logsByStage