From b74e040b0d84afe4dced23ab73bcf4c9a7aef21d Mon Sep 17 00:00:00 2001
From: Kefu Chai <kchai@redhat.com>
Date: Mon, 22 Feb 2021 11:28:07 +0800
Subject: [PATCH] mgr/influx: define commands uing CLICommand

Signed-off-by: Kefu Chai <kchai@redhat.com>
---
 src/pybind/mgr/influx/module.py   | 48 +++++++++++++++++--------------
 src/pybind/mgr/selftest/module.py |  6 ++--
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/src/pybind/mgr/influx/module.py b/src/pybind/mgr/influx/module.py
index 56b53e0bb8a..ab08c865230 100644
--- a/src/pybind/mgr/influx/module.py
+++ b/src/pybind/mgr/influx/module.py
@@ -6,8 +6,9 @@ import queue
 import json
 import errno
 import time
+from typing import Tuple
 
-from mgr_module import MgrModule, Option
+from mgr_module import CLICommand, CLIReadCommand, CLIWriteCommand, MgrModule, Option
 
 try:
     from influxdb import InfluxDBClient
@@ -418,28 +419,33 @@ class Module(MgrModule):
 
         return json.dumps(result, indent=2, sort_keys=True)
 
-    def handle_command(self, inbuf, cmd):
-        if cmd['prefix'] == 'influx config-show':
-            return 0, json.dumps(self.config, sort_keys=True), ''
-        elif cmd['prefix'] == 'influx config-set':
-            key = cmd['key']
-            value = cmd['value']
-            if not value:
-                return -errno.EINVAL, '', 'Value should not be empty or None'
+    @CLIReadCommand('influx config-show')
+    def config_show(self) -> Tuple[int, str, str]:
+        """
+        Show current configuration
+        """
+        return 0, json.dumps(self.config, sort_keys=True), ''
 
-            self.log.debug('Setting configuration option %s to %s', key, value)
-            try:
-                self.set_module_option(key, value)
-                self.set_config_option(key, value)
-                return 0, 'Configuration option {0} updated'.format(key), ''
-            except ValueError as e:
-                return -errno.EINVAL, '', str(e)
-        elif cmd['prefix'] == 'influx send':
-            self.send_to_influx()
-            return 0, 'Sending data to Influx', ''
+    @CLIWriteCommand('influx config-set')
+    def config_set(self, key: str, value: str) -> Tuple[int, str, str]:
+        if not value:
+            return -errno.EINVAL, '', 'Value should not be empty'
 
-        return (-errno.EINVAL, '',
-                "Command not found '{0}'".format(cmd['prefix']))
+        self.log.debug('Setting configuration option %s to %s', key, value)
+        try:
+            self.set_module_option(key, value)
+            self.set_config_option(key, value)
+            return 0, 'Configuration option {0} updated'.format(key), ''
+        except ValueError as e:
+            return -errno.EINVAL, '', str(e)
+
+    @CLICommand('influx send')
+    def send(self) -> Tuple[int, str, str]:
+        """
+        Force sending data to Influx
+        """
+        self.send_to_influx()
+        return 0, 'Sending data to Influx', ''
 
     def serve(self):
         if InfluxDBClient is None:
diff --git a/src/pybind/mgr/selftest/module.py b/src/pybind/mgr/selftest/module.py
index cf354000018..dd0c381671d 100644
--- a/src/pybind/mgr/selftest/module.py
+++ b/src/pybind/mgr/selftest/module.py
@@ -402,7 +402,7 @@ class Module(MgrModule):
 
     def _test_remote_calls(self):
         # Test making valid call
-        self.remote("influx", "handle_command", "", {"prefix": "influx self-test"})
+        self.remote("influx", "self_test")
 
         # Test calling module that exists but isn't enabled
         # (arbitrarily pick a non-always-on module to use)
@@ -423,7 +423,7 @@ class Module(MgrModule):
 
         # Test calling module that doesn't exist
         try:
-            self.remote("idontexist", "handle_command", {"prefix": "influx self-test"})
+            self.remote("idontexist", "self_test")
         except ImportError:
             pass
         else:
@@ -431,7 +431,7 @@ class Module(MgrModule):
 
         # Test calling method that doesn't exist
         try:
-            self.remote("influx", "idontexist", {"prefix": "influx self-test"})
+            self.remote("influx", "idontexist")
         except NameError:
             pass
         else:
-- 
2.39.5