From c729e282c505ae60e4215d35a209b866356209d0 Mon Sep 17 00:00:00 2001
From: bnechyporenko <bnechyporenko@bol.com>
Date: Mon, 19 Feb 2024 19:40:38 +0100
Subject: [PATCH] wip

Signed-off-by: bnechyporenko <bnechyporenko@bol.com>
---
 package.json                                  |  1 -
 .../tasks/DatabaseTaskStore.test.ts           | 16 ++++++-----
 .../tasks/StorageTaskBroker.test.ts           | 27 +++++++++++++++++++
 .../src/scaffolder/tasks/StorageTaskBroker.ts | 17 ++++++++++--
 yarn.lock                                     |  1 -
 5 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/package.json b/package.json
index e224d1e58a..afc86f2460 100644
--- a/package.json
+++ b/package.json
@@ -76,7 +76,6 @@
     "@backstage/repo-tools": "workspace:*",
     "@changesets/cli": "^2.14.0",
     "@octokit/rest": "^19.0.3",
-    "@playwright/test": "^1.32.3",
     "@spotify/eslint-plugin": "^14.1.3",
     "@spotify/prettier-config": "^14.0.0",
     "@techdocs/cli": "workspace:*",
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.test.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.test.ts
index 8fda4b5090..fe9cf3c661 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.test.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.test.ts
@@ -199,9 +199,11 @@ describe('DatabaseTaskStore', () => {
     await store.saveTaskState({
       taskId,
       state: {
-        'repo.create': {
-          status: 'success',
-          value: { repoUrl: 'https://github.com/backstage/backstage.git' },
+        checkpoints: {
+          'repo.create': {
+            status: 'success',
+            value: { repoUrl: 'https://github.com/backstage/backstage.git' },
+          },
         },
       },
     });
@@ -210,9 +212,11 @@ describe('DatabaseTaskStore', () => {
 
     expect(state).toStrictEqual({
       state: {
-        'repo.create': {
-          status: 'success',
-          value: { repoUrl: 'https://github.com/backstage/backstage.git' },
+        checkpoints: {
+          'repo.create': {
+            status: 'success',
+            value: { repoUrl: 'https://github.com/backstage/backstage.git' },
+          },
         },
       },
     });
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.test.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.test.ts
index 3d2fb28538..f4bae43dbf 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.test.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.test.ts
@@ -242,4 +242,31 @@ describe('StorageTaskBroker', () => {
     const promise = broker.list({ createdBy: 'user:default/foo' });
     await expect(promise).resolves.toEqual({ tasks: [task] });
   });
+
+  it('should handle checkpoints in task state', async () => {
+    const broker = new StorageTaskBroker(storage, logger);
+
+    await broker.dispatch({
+      spec: { steps: [] } as unknown as TaskSpec,
+      createdBy: 'user:default/foo',
+    });
+
+    const taskA = await broker.claim();
+    await taskA.updateCheckpoint?.({
+      key: 'repo.create',
+      status: 'success',
+      value: 'https://github.com/backstage/backstage.git',
+    });
+
+    expect(await taskA.getTaskState?.()).toEqual({
+      state: {
+        checkpoints: {
+          'repo.create': {
+            status: 'success',
+            value: 'https://github.com/backstage/backstage.git',
+          },
+        },
+      },
+    });
+  });
 });
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
index c437795778..fec745a653 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
@@ -31,6 +31,19 @@ import {
 import { TaskStore } from './types';
 import { readDuration } from './helper';
 
+type TaskState = {
+  checkpoints: {
+    [key: string]:
+      | {
+          status: 'failed';
+          reason: string;
+        }
+      | {
+          status: 'success';
+          value: JsonValue;
+        };
+  };
+};
 /**
  * TaskManager
  *
@@ -115,9 +128,9 @@ export class TaskManager implements TaskContext {
   ): Promise<void> {
     const { key, ...value } = options;
     if (this.task.state) {
-      this.task.state[key] = value;
+      (this.task.state as TaskState).checkpoints[key] = value;
     } else {
-      this.task.state = { [key]: value };
+      this.task.state = { checkpoints: { [key]: value } };
     }
     await this.storage.saveTaskState?.({
       taskId: this.task.taskId,
diff --git a/yarn.lock b/yarn.lock
index 65a4a0a508..cd5f483ddc 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -41210,7 +41210,6 @@ __metadata:
     "@changesets/cli": ^2.14.0
     "@manypkg/get-packages": ^1.1.3
     "@octokit/rest": ^19.0.3
-    "@playwright/test": ^1.32.3
     "@spotify/eslint-plugin": ^14.1.3
     "@spotify/prettier-config": ^14.0.0
     "@techdocs/cli": "workspace:*"