From 02b0b7b9f99b85de34e56168cf4ccb6955f4c553 Mon Sep 17 00:00:00 2001
From: Elliott Marquez <5981958+e111077@users.noreply.github.com>
Date: Thu, 29 Sep 2022 11:35:30 -0700
Subject: [PATCH] Task add onComplete and onError (#3287)

---
 .changeset/lazy-moons-train.md           |  5 ++
 packages/labs/task/src/task.ts           | 16 +++++
 packages/labs/task/src/test/task_test.ts | 76 +++++++++++++++++++++++-
 3 files changed, 94 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/lazy-moons-train.md

diff --git a/.changeset/lazy-moons-train.md b/.changeset/lazy-moons-train.md
new file mode 100644
index 0000000000..885a5db669
--- /dev/null
+++ b/.changeset/lazy-moons-train.md
@@ -0,0 +1,5 @@
+---
+'@lit-labs/task': minor
+---
+
+Adds onComplete and onError callbacks
diff --git a/packages/labs/task/src/task.ts b/packages/labs/task/src/task.ts
index 0f1aa19b28..cf12672cb7 100644
--- a/packages/labs/task/src/task.ts
+++ b/packages/labs/task/src/task.ts
@@ -44,6 +44,8 @@ export interface TaskConfig<T extends ReadonlyArray<unknown>, R> {
   task: TaskFunction<T, R>;
   args?: ArgsFunction<T>;
   autoRun?: boolean;
+  onComplete?: (value: R) => unknown;
+  onError?: (error: unknown) => unknown;
 }
 
 // TODO(sorvell): Some issues:
@@ -109,6 +111,8 @@ export class Task<
   private _host: ReactiveControllerHost;
   private _value?: R;
   private _error?: unknown;
+  private _onComplete?: (result: R) => unknown;
+  private _onError?: (error: unknown) => unknown;
   status: TaskStatus = TaskStatus.INITIAL;
 
   /**
@@ -144,6 +148,8 @@ export class Task<
       typeof task === 'object' ? task : ({task, args} as TaskConfig<T, R>);
     this._task = taskConfig.task;
     this._getArgs = taskConfig.args;
+    this._onComplete = taskConfig.onComplete;
+    this._onError = taskConfig.onError;
     if (taskConfig.autoRun !== undefined) {
       this.autoRun = taskConfig.autoRun;
     }
@@ -211,9 +217,19 @@ export class Task<
         this.status = TaskStatus.INITIAL;
       } else {
         if (error === undefined) {
+          try {
+            this._onComplete?.(result as R);
+          } catch {
+            // Ignore user errors from onComplete.
+          }
           this.status = TaskStatus.COMPLETE;
           this._resolveTaskComplete(result as R);
         } else {
+          try {
+            this._onError?.(error);
+          } catch {
+            // Ignore user errors from onError.
+          }
           this.status = TaskStatus.ERROR;
           this._rejectTaskComplete(error);
         }
diff --git a/packages/labs/task/src/test/task_test.ts b/packages/labs/task/src/test/task_test.ts
index ab3960b2ea..848ed77134 100644
--- a/packages/labs/task/src/test/task_test.ts
+++ b/packages/labs/task/src/test/task_test.ts
@@ -27,7 +27,7 @@ suite('Task', () => {
     b: string;
     c?: string;
     resolveTask: () => void;
-    rejectTask: () => void;
+    rejectTask: (error?: string) => void;
     taskValue?: string;
     renderedStatus?: string;
   }
@@ -46,7 +46,7 @@ suite('Task', () => {
       c?: string;
 
       resolveTask!: () => void;
-      rejectTask!: () => void;
+      rejectTask!: (error?: string) => void;
 
       taskValue?: string;
       renderedStatus?: string;
@@ -56,7 +56,7 @@ suite('Task', () => {
         const taskConfig = {
           task: (...args: unknown[]) =>
             new Promise((resolve, reject) => {
-              this.rejectTask = () => reject(`error`);
+              this.rejectTask = (error = 'error') => reject(error);
               this.resolveTask = () => resolve(args.join(','));
             }),
         };
@@ -458,4 +458,74 @@ suite('Task', () => {
       );
     };
   });
+
+  test('onComplete callback is called', async () => {
+    let numOnCompleteInvocations = 0;
+    let lastOnCompleteResult: string | undefined = undefined;
+    const el = getTestElement({
+      args: () => [el.a, el.b],
+      onComplete: (result) => {
+        numOnCompleteInvocations++;
+        lastOnCompleteResult = result;
+      },
+    });
+    await renderElement(el);
+    assert.equal(el.task.status, TaskStatus.PENDING);
+    assert.equal(numOnCompleteInvocations, 0);
+    assert.equal(lastOnCompleteResult, undefined);
+    el.resolveTask();
+    await tasksUpdateComplete();
+    assert.equal(el.task.status, TaskStatus.COMPLETE);
+    assert.equal(numOnCompleteInvocations, 1);
+    assert.equal(lastOnCompleteResult, 'a,b');
+
+    numOnCompleteInvocations = 0;
+
+    // Called after every task completion.
+    el.a = 'a1';
+    await tasksUpdateComplete();
+    assert.equal(el.task.status, TaskStatus.PENDING);
+    assert.equal(numOnCompleteInvocations, 0);
+    assert.equal(lastOnCompleteResult, 'a,b');
+    el.resolveTask();
+    await tasksUpdateComplete();
+    assert.equal(el.task.status, TaskStatus.COMPLETE);
+    assert.equal(numOnCompleteInvocations, 1);
+    assert.equal(lastOnCompleteResult, 'a1,b');
+  });
+
+  test('onError callback is called', async () => {
+    let numOnErrorInvocations = 0;
+    let lastOnErrorResult: string | undefined = undefined;
+    const el = getTestElement({
+      args: () => [el.a, el.b],
+      onError: (error) => {
+        numOnErrorInvocations++;
+        lastOnErrorResult = error as string;
+      },
+    });
+    await renderElement(el);
+    assert.equal(el.task.status, TaskStatus.PENDING);
+    assert.equal(numOnErrorInvocations, 0);
+    assert.equal(lastOnErrorResult, undefined);
+    el.rejectTask('error');
+    await tasksUpdateComplete();
+    assert.equal(el.task.status, TaskStatus.ERROR);
+    assert.equal(numOnErrorInvocations, 1);
+    assert.equal(lastOnErrorResult, 'error');
+
+    numOnErrorInvocations = 0;
+
+    // Called after every task error.
+    el.a = 'a1';
+    await tasksUpdateComplete();
+    assert.equal(el.task.status, TaskStatus.PENDING);
+    assert.equal(numOnErrorInvocations, 0);
+    assert.equal(lastOnErrorResult, 'error');
+    el.rejectTask('error2');
+    await tasksUpdateComplete();
+    assert.equal(el.task.status, TaskStatus.ERROR);
+    assert.equal(numOnErrorInvocations, 1);
+    assert.equal(lastOnErrorResult, 'error2');
+  });
 });