feat(cil): support for reading messages from files in bench

emqx · Apr 23, 2024 · a6cb136 · a6cb136
1 parent 25412b0
commit a6cb136
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 4 deletions.
diff --git a/cli/src/index.ts b/cli/src/index.ts
@@ -692,6 +692,10 @@ export class Commander {
         '--config [PATH]',
         'load the parameters from the local configuration file, which supports json and yaml format, default path is ./mqttx-cli-config.json',
       )
+      .option(
+        '--file-read <PATH>',
+        'read the message body from the file',
+      )
       .allowUnknownOption(false)
       .action(simulatePub)
 

diff --git a/cli/src/lib/pub.ts b/cli/src/lib/pub.ts
@@ -173,13 +173,13 @@ const pub = (options: PublishOptions) => {
     }
 
     try {
-      signale.pending('Reading file...')
+      basicLog.fileReading()
       const bufferData = readFile(filePath)
       if (bufferData.length >= MQTT_SINGLE_MESSAGE_BYTE_LIMIT) {
         signale.error('File size over 256MB not supported by MQTT.')
         process.exit(1)
       }
-      signale.success('File read successfully')
+      basicLog.fileReadSuccess()
 
       pubOpts.message = bufferData
       send(config, connOpts, pubOpts)
@@ -237,10 +237,37 @@ const multiPub = async (commandType: CommandType, options: BenchPublishOptions |
     topic,
     clientId,
     message,
+    fileRead,
     verbose,
     maximumReconnectTimes,
   } = options
 
+  const handleFileRead = () => {
+    const filePath = processPath(fileRead!)
+    if (!filePath) {
+      signale.error('File path is required when reading from file.')
+      process.exit(1)
+    }
+
+    try {
+      basicLog.fileReading()
+      const bufferData = readFile(filePath)
+      if (bufferData.length >= MQTT_SINGLE_MESSAGE_BYTE_LIMIT) {
+        signale.error('File size over 256MB not supported by MQTT.')
+        process.exit(1)
+      }
+      basicLog.fileReadSuccess()
+      return bufferData
+    } catch(error) {
+      signale.error('Failed to read file:', error)
+      process.exit(1)
+    }
+  }
+  let fileData: Buffer | string
+  if(fileRead) {
+    fileData = handleFileRead()
+  }
+
   checkTopicExists(topic, commandType)
 
   const connOpts = parseConnectOptions(options, 'pub')
@@ -327,6 +354,9 @@ const multiPub = async (commandType: CommandType, options: BenchPublishOptions |
               }
               publishMessage = simulationResult.message
             }
+            if(fileRead) {
+              publishMessage = fileData
+            }
             client.publish(publishTopic, publishMessage, pubOpts.opts, (err) => {
               inFlightMessageCount -= 1
               if (err) {

diff --git a/cli/src/types/global.d.ts b/cli/src/types/global.d.ts
@@ -118,7 +118,7 @@ declare global {
 
   type OmitSubscribeOptions = Omit<
     SubscribeOptions,
-    'format' | 'outputMode' | 'protobufPath' | 'protobufMessageName' | 'debug'
+    'format' | 'outputMode' | 'protobufPath' | 'protobufMessageName' | 'debug' | 'fileWrite' | 'fileSave'
   >
 
   interface BenchSubscribeOptions extends OmitSubscribeOptions {

diff --git a/cli/src/utils/signale.ts b/cli/src/utils/signale.ts
@@ -52,7 +52,9 @@ const basicLog = {
     const { reasonCode } = packet
     const reason = reasonCode === 0 ? 'Normal disconnection' : getErrorReason(reasonCode)
     signale.warn(`${clientId ? `Client ID: ${clientId}, ` : ''}The Broker has actively disconnected, Reason: ${reason} (Code: ${reasonCode})`)
-  }
+  },
+  fileReading: () => signale.await('Reading file...'),
+  fileReadSuccess: () => signale.success('Read file successfully'),
 }
 
 const benchLog = {