From c7aa19e904891339bfba957adbd2eabbcb7a47dd Mon Sep 17 00:00:00 2001
From: ajmnz <arnau@ajb.cat>
Date: Thu, 27 Jun 2024 14:43:35 +0200
Subject: [PATCH] allow awaiting toast.promise

---
 src/state.ts             | 52 +++++++++++++++++++++++++---------------
 test/tests/basic.spec.ts | 12 ++++++++++
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/src/state.ts b/src/state.ts
index 67260e7..ff365de 100644
--- a/src/state.ts
+++ b/src/state.ts
@@ -124,26 +124,30 @@ class Observer {
     const p = promise instanceof Promise ? promise : promise();
 
     let shouldDismiss = id !== undefined;
+    let result: ['resolve', ToastData] | ['reject', unknown];
 
-    p.then(async (response) => {
-      if (isHttpResponse(response) && !response.ok) {
-        shouldDismiss = false;
-        const message =
-          typeof data.error === 'function' ? await data.error(`HTTP error! status: ${response.status}`) : data.error;
-        const description =
-          typeof data.description === 'function'
-            ? await data.description(`HTTP error! status: ${response.status}`)
-            : data.description;
-        this.create({ id, type: 'error', message, description });
-      } else if (data.success !== undefined) {
-        shouldDismiss = false;
-        const message = typeof data.success === 'function' ? await data.success(response) : data.success;
-        const description =
-          typeof data.description === 'function' ? await data.description(response) : data.description;
-        this.create({ id, type: 'success', message, description });
-      }
-    })
+    const originalPromise = p
+      .then(async (response) => {
+        result = ['resolve', response];
+        if (isHttpResponse(response) && !response.ok) {
+          shouldDismiss = false;
+          const message =
+            typeof data.error === 'function' ? await data.error(`HTTP error! status: ${response.status}`) : data.error;
+          const description =
+            typeof data.description === 'function'
+              ? await data.description(`HTTP error! status: ${response.status}`)
+              : data.description;
+          this.create({ id, type: 'error', message, description });
+        } else if (data.success !== undefined) {
+          shouldDismiss = false;
+          const message = typeof data.success === 'function' ? await data.success(response) : data.success;
+          const description =
+            typeof data.description === 'function' ? await data.description(response) : data.description;
+          this.create({ id, type: 'success', message, description });
+        }
+      })
       .catch(async (error) => {
+        result = ['reject', error];
         if (data.error !== undefined) {
           shouldDismiss = false;
           const message = typeof data.error === 'function' ? await data.error(error) : data.error;
@@ -161,7 +165,17 @@ class Observer {
         data.finally?.();
       });
 
-    return id;
+    const unwrap = () =>
+      new Promise<ToastData>((resolve, reject) =>
+        originalPromise.then(() => (result[0] === 'reject' ? reject(result[1]) : resolve(result[1]))).catch(reject),
+      );
+
+    if (typeof id !== 'string' && typeof id !== 'number') {
+      // cannot Object.assign on undefined
+      return { unwrap };
+    } else {
+      return Object.assign(id, { unwrap });
+    }
   };
 
   custom = (jsx: (id: number | string) => React.ReactElement, data?: ExternalToast) => {
diff --git a/test/tests/basic.spec.ts b/test/tests/basic.spec.ts
index 80daadb..c40c712 100644
--- a/test/tests/basic.spec.ts
+++ b/test/tests/basic.spec.ts
@@ -1,4 +1,5 @@
 import { expect, test } from '@playwright/test';
+import { toast } from 'sonner';
 
 test.beforeEach(async ({ page }) => {
   await page.goto('/');
@@ -28,6 +29,17 @@ test.describe('Basic functionality', () => {
     await expect(page.getByText('Loaded')).toHaveCount(1);
   });
 
+  test('handle toast promise rejections', async ({ page }) => {
+    const rejectedPromise = new Promise((_, reject) => setTimeout(() => reject(new Error('Promise rejected')), 100));
+    try {
+      toast.promise(rejectedPromise, {});
+    } catch {
+      throw new Error('Promise should not have rejected without unwrap');
+    }
+
+    await expect(toast.promise(rejectedPromise, {}).unwrap()).rejects.toThrow('Promise rejected');
+  });
+
   test('render custom jsx in toast', async ({ page }) => {
     await page.getByTestId('custom').click();
     await expect(page.getByText('jsx')).toHaveCount(1);