fix(cli): code optimization

emqx · May 9, 2024 · afb2992 · afb2992
1 parent 98b3c1d
commit afb2992
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 16 deletions.
diff --git a/cli/src/lib/pub.ts b/cli/src/lib/pub.ts
@@ -10,7 +10,7 @@ import delay from '../utils/delay'
 import { saveConfig, loadConfig } from '../utils/config'
 import { loadSimulator } from '../utils/simulate'
 import { serializeProtobufToBuffer } from '../utils/protobuf'
-import { readFile, processPath, fileDataSplitter } from '../utils/fileUtils'
+import { readFile, processPath, fileDataSplitter, getPublishMessageFromFile } from '../utils/fileUtils'
 import convertPayload from '../utils/convertPayload'
 import * as Debug from 'debug'
 import _ from 'lodash'
@@ -288,10 +288,8 @@ const multiPub = async (commandType: CommandType, options: BenchPublishOptions |
 
   for (let i = 1; i <= count; i++) {
     // Duplicate splited messages array for each connection
-    let dupSplitedMessageArr: string[] = []
-    if (split && fileRead && splitedMessageArr.length !== 0) {
-      dupSplitedMessageArr = _.cloneDeep(splitedMessageArr)
-    }
+    const dupSplitedMessageArr: string[] =
+      split && fileRead && splitedMessageArr.length !== 0 ? _.cloneDeep(splitedMessageArr) : []
 
     ;((i: number, connOpts: mqtt.IClientOptions) => {
       const opts = { ...connOpts }
@@ -336,17 +334,10 @@ const multiPub = async (commandType: CommandType, options: BenchPublishOptions |
               publishMessage = simulationResult.message
             }
             if (fileRead) {
-              if (split) {
-                if (dupSplitedMessageArr.length === 0) {
-                  await delay(1000)
-                  signale.success(`All ${total} messages from the ${fileRead} have been successfully sent.`)
-                  process.exit(0)
-                }
-                const unshiftedMessage = dupSplitedMessageArr.shift()
-                publishMessage = Buffer.from(unshiftedMessage!)
-              } else {
-                publishMessage = fileData
-              }
+              publishMessage = await getPublishMessageFromFile(split, dupSplitedMessageArr, fileData, {
+                total,
+                fileRead,
+              })
             }
             client.publish(publishTopic, publishMessage, pubOpts.opts, (err) => {
               inFlightMessageCount -= 1

diff --git a/cli/src/utils/fileUtils.ts b/cli/src/utils/fileUtils.ts
@@ -2,6 +2,7 @@ import fs from 'fs'
 import path from 'path'
 import YAML from 'js-yaml'
 import signale from 'signale'
+import delay from './delay'
 
 const processPath = (savePath: boolean | string, defaultPath?: string) => {
   let filePath = ''
@@ -97,6 +98,25 @@ const fileDataSplitter = (data: string | Buffer, split: string): string[] => {
   return stringData.split(splitRegex)
 }
 
+const getPublishMessageFromFile = async (
+  split: string,
+  dupSplitedMessageArr: string[],
+  fileData: string | Buffer,
+  meta: { total: number; fileRead: string },
+): Promise<Buffer | string> => {
+  if (!split) {
+    return fileData
+  }
+  if (dupSplitedMessageArr.length === 0) {
+    await delay(1000)
+    signale.success(`All ${meta.total} messages from the ${meta.fileRead} have been successfully sent.`)
+    process.exit(0)
+  }
+
+  const unshiftedMessage = dupSplitedMessageArr.shift()
+  return Buffer.from(unshiftedMessage!)
+}
+
 export {
   processPath,
   getPathExtname,
@@ -109,4 +129,5 @@ export {
   appendFile,
   createNextNumberedFileName,
   fileDataSplitter,
+  getPublishMessageFromFile,
 }