From 1751fd62ca893797547b0d20529beea756642816 Mon Sep 17 00:00:00 2001
From: starriv <starriv.993@gmail.com>
Date: Tue, 10 May 2022 19:44:52 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9E=20fix:=20return=20method=20for=20a?=
 =?UTF-8?q?pi=20with=20task=20type?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/__tests__/promisify.test.ts           | 15 +++++++
 .../remax-framework-shared/src/promisify.ts   | 40 +++++++++++++++----
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/packages/remax-framework-shared/src/__tests__/promisify.test.ts b/packages/remax-framework-shared/src/__tests__/promisify.test.ts
index 77e1086f8..b9421344f 100644
--- a/packages/remax-framework-shared/src/__tests__/promisify.test.ts
+++ b/packages/remax-framework-shared/src/__tests__/promisify.test.ts
@@ -27,6 +27,21 @@ describe('promisify', () => {
     expect(success).toBeCalledTimes(1);
   });
 
+  it('resolve by complete param', async () => {
+    function api({ complete }: any) {
+      complete('complete');
+    }
+
+    const promisifyAPI = promisify(api);
+    const complete = jest.fn();
+
+    await promisifyAPI({
+      complete,
+    });
+
+    expect(complete).toBeCalledTimes(1);
+  });
+
   it('promisify api and reject when failed', () => {
     function api({ fail }: any) {
       fail('error');
diff --git a/packages/remax-framework-shared/src/promisify.ts b/packages/remax-framework-shared/src/promisify.ts
index 53d3ecf7a..df3a451dd 100644
--- a/packages/remax-framework-shared/src/promisify.ts
+++ b/packages/remax-framework-shared/src/promisify.ts
@@ -1,16 +1,17 @@
-export interface PromisifyArgs<SuccessArg, FailArg> {
+export interface PromisifyArgs<SuccessArg, FailArg, CompleteArg> {
   success?: (args: SuccessArg) => void;
   fail?: (args: FailArg) => void;
+  complete?: (args: CompleteArg) => void;
 }
 
-export function promisify<Arg = any, SuccessArg = any, FailArg = any>(
-  api: (arg: Arg & PromisifyArgs<SuccessArg, FailArg>) => void
+export function promisify<Arg = any, SuccessArg = any, FailArg = any, CompleteArg = any>(
+  api: (arg: Arg & PromisifyArgs<SuccessArg, FailArg, CompleteArg>) => any
 ) {
-  return (arg: Arg & PromisifyArgs<SuccessArg, FailArg> = {} as Arg) => {
-    return new Promise<SuccessArg>((resolve, reject) => {
+  return (arg: Arg & PromisifyArgs<SuccessArg, FailArg, CompleteArg> = {} as Arg) => {
+    let task: any = null;
+    const p = new Promise<{ [key: string]: any }>((resolve, reject) => {
       const promisifyArg: any = arg;
-
-      api({
+      task = api({
         ...promisifyArg,
         success: (res: SuccessArg) => {
           if (promisifyArg && typeof promisifyArg.success === 'function') {
@@ -24,7 +25,32 @@ export function promisify<Arg = any, SuccessArg = any, FailArg = any>(
           }
           reject(res);
         },
+        complete: (res: CompleteArg) => {
+          if (promisifyArg && typeof promisifyArg.complete === 'function') {
+            promisifyArg.complete(res);
+          }
+          resolve(res);
+        },
       });
     });
+
+    const taskMethods = [
+      'abort',
+      'onHeadersReceived',
+      'offHeadersReceived',
+      'onProgressUpdate',
+      'offProgressUpdate',
+      'onChunkReceived',
+      'offChunkReceived',
+    ];
+    task &&
+      taskMethods.forEach(method => {
+        if (method in task) {
+          // @ts-ignore
+          p[method] = task[method].bind(task);
+        }
+      });
+
+    return p;
   };
 }