1 files changed, 13 insertions, 12 deletions
diff --git a/src/pump_bot.py b/src/pump_bot.py
index 48efec4..ab73097 100755
--- a/src/pump_bot.py
+++ b/src/pump_bot.py
@@ -10,16 +10,17 @@ from home.telegram import bot
 from home.telegram._botutil import user_any_name
 from home.relay.sunxi_h3_client import RelayClient
 from home.api.types import BotType
-from home.mqtt import MqttNode, MqttModule, MqttPayload, add_mqtt_module
-from home.mqtt.module.relay import MqttPowerStatusPayload
+from home.mqtt import MqttNode, MqttWrapper, MqttPayload
+from home.mqtt.module.relay import MqttPowerStatusPayload, MqttRelayModule
 from home.mqtt.module.temphum import MqttTemphumDataPayload
 from home.mqtt.module.diagnostics import InitialDiagnosticsPayload, DiagnosticsPayload
 
 
 config.load('pump_bot')
 
-mqtt: Optional[MqttNode] = None
-mqtt_relay_module: Optional[MqttModule] = None
+mqtt: Optional[MqttWrapper] = None
+mqtt_node: Optional[MqttNode] = None
+mqtt_relay_module: Optional[MqttRelayModule] = None
 time_format = '%d.%m.%Y, %H:%M:%S'
 
 watering_mcu_status = {
@@ -117,13 +118,13 @@ def off(ctx: bot.Context, silent=False) -> None:
 
 
 def watering_on(ctx: bot.Context) -> None:
-    mqtt_relay_module.switchpower(mqtt, True, config.get('mqtt_water_relay.secret'))
+    mqtt_relay_module.switchpower(True, config.get('mqtt_water_relay.secret'))
     ctx.reply(ctx.lang('sent'))
     notify(ctx.user, UserAction.WATERING_ON)
 
 
 def watering_off(ctx: bot.Context) -> None:
-    mqtt_relay_module.switchpower(mqtt, False, config.get('mqtt_water_relay.secret'))
+    mqtt_relay_module.switchpower(False, config.get('mqtt_water_relay.secret'))
     ctx.reply(ctx.lang('sent'))
     notify(ctx.user, UserAction.WATERING_OFF)
 
@@ -233,14 +234,15 @@ def mqtt_payload_callback(payload: MqttPayload):
 
 
 if __name__ == '__main__':
-    mqtt = MqttNode(node_id=config.get('mqtt_water_relay.node_id'))
+    mqtt = MqttWrapper()
+    mqtt_node = MqttNode(node_id=config.get('mqtt_water_relay.node_id'))
     if is_development_mode():
-        add_mqtt_module(mqtt, 'diagnostics')
+        mqtt_node.load_module('diagnostics')
 
-    mqtt_relay_module = add_mqtt_module(mqtt, 'temphum')
-    mqtt_relay_module = add_mqtt_module(mqtt, 'relay')
+    mqtt_node.load_module('temphum')
+    mqtt_relay_module = mqtt_node.load_module('relay')
 
-    mqtt.add_payload_callback(mqtt_payload_callback)
+    mqtt_node.add_payload_callback(mqtt_payload_callback)
 
     mqtt.configure_tls()
     mqtt.connect_and_loop(loop_forever=False)
@@ -252,4 +254,3 @@ if __name__ == '__main__':
         mqtt.disconnect()
     except:
         pass
-