From 812310a816174d0a0d013fcb7b50787153f9a6b2 Mon Sep 17 00:00:00 2001
From: Danish Arora <35004822+danisharora099@users.noreply.github.com>
Date: Tue, 22 Aug 2023 15:11:34 +0530
Subject: [PATCH] chore: refactor LightPush `send` (#1487)

* refactor lightpush send

* add trycatch
---
 packages/core/src/lib/light_push/index.ts | 71 +++++++++++++++++------
 1 file changed, 53 insertions(+), 18 deletions(-)

diff --git a/packages/core/src/lib/light_push/index.ts b/packages/core/src/lib/light_push/index.ts
index bca6cd9ea9..a0259b3d12 100644
--- a/packages/core/src/lib/light_push/index.ts
+++ b/packages/core/src/lib/light_push/index.ts
@@ -38,37 +38,70 @@ class LightPush extends BaseProtocol implements ILightPush {
     this.options = options || {};
   }
 
+  private async preparePushMessage(
+    encoder: IEncoder,
+    message: IMessage,
+    pubSubTopic: string
+  ): Promise<{
+    query: PushRpc | null;
+    error?: SendError;
+  }> {
+    if (!isSizeValid(message.payload)) {
+      log("Failed to send waku light push: message is bigger than 1MB");
+      return { query: null, error: SendError.SIZE_TOO_BIG };
+    }
+
+    const protoMessage = await encoder.toProtoObj(message);
+    if (!protoMessage) {
+      log("Failed to encode to protoMessage, aborting push");
+      return {
+        query: null,
+        error: SendError.ENCODE_FAILED
+      };
+    }
+
+    const query = PushRpc.createRequest(protoMessage, pubSubTopic);
+    return { query };
+  }
+
   async send(
     encoder: IEncoder,
     message: IMessage,
     opts?: ProtocolOptions
   ): Promise<SendResult> {
     const { pubSubTopic = DefaultPubSubTopic } = this.options;
-
-    const peer = await this.getPeer(opts?.peerId);
-    const stream = await this.newStream(peer);
-
     const recipients: PeerId[] = [];
     let error: undefined | SendError = undefined;
 
+    let query: PushRpc | null = null;
+
     try {
-      if (!isSizeValid(message.payload)) {
-        log("Failed to send waku light push: message is bigger that 1MB");
-        return {
-          recipients,
-          error: SendError.SIZE_TOO_BIG
-        };
-      }
+      const { query: preparedQuery, error: preparationError } =
+        await this.preparePushMessage(encoder, message, pubSubTopic);
 
-      const protoMessage = await encoder.toProtoObj(message);
-      if (!protoMessage) {
-        log("Failed to encode to protoMessage, aborting push");
+      if (preparationError) {
         return {
           recipients,
-          error: SendError.ENCODE_FAILED
+          error: preparationError
         };
       }
-      const query = PushRpc.createRequest(protoMessage, pubSubTopic);
+
+      query = preparedQuery;
+    } catch (error) {
+      log("Failed to prepare push message", error);
+    }
+
+    if (!query) {
+      return {
+        recipients,
+        error: SendError.GENERIC_FAIL
+      };
+    }
+
+    const peer = await this.getPeer(opts?.peerId);
+    const stream = await this.newStream(peer);
+
+    try {
       const res = await pipe(
         [query.encode()],
         lp.encode,
@@ -76,6 +109,7 @@ class LightPush extends BaseProtocol implements ILightPush {
         lp.decode,
         async (source) => await all(source)
       );
+
       try {
         const bytes = new Uint8ArrayList();
         res.forEach((chunk) => {
@@ -98,9 +132,10 @@ class LightPush extends BaseProtocol implements ILightPush {
       log("Failed to send waku light push request", err);
       error = SendError.GENERIC_FAIL;
     }
+
     return {
-      error,
-      recipients
+      recipients,
+      error
     };
   }
 }