1 files changed, 11 insertions, 11 deletions
diff --git a/src/pump_mqtt_bot.py b/src/pump_mqtt_bot.py
index accafcb..d3b6de4 100755
--- a/src/pump_mqtt_bot.py
+++ b/src/pump_mqtt_bot.py
@@ -8,10 +8,10 @@ from telegram import ReplyKeyboardMarkup, User
 from home.config import config
 from home.telegram import bot
 from home.telegram._botutil import user_any_name
-from home.api.types import BotType
-from home.mqtt import MQTTRelay, MQTTRelayState, MQTTRelayDevice
-from home.mqtt.payload import MQTTPayload
-from home.mqtt.payload.relay import InitialStatPayload, StatPayload
+from home.mqtt.esp import MqttEspDevice
+from home.mqtt import MqttRelay, MqttRelayState
+from home.mqtt.payload import MqttPayload
+from home.mqtt.payload.relay import InitialDiagnosticsPayload, DiagnosticsPayload
 
 
 config.load('pump_mqtt_bot')
@@ -70,8 +70,8 @@ bot.lang.en(
 )
 
 
-mqtt_relay: Optional[MQTTRelay] = None
-relay_state = MQTTRelayState()
+mqtt_relay: Optional[MqttRelay] = None
+relay_state = MqttRelayState()
 
 
 class UserAction(Enum):
@@ -79,10 +79,10 @@ class UserAction(Enum):
     OFF = 'off'
 
 
-def on_mqtt_message(home_id, message: MQTTPayload):
-    if isinstance(message, InitialStatPayload) or isinstance(message, StatPayload):
+def on_mqtt_message(home_id, message: MqttPayload):
+    if isinstance(message, InitialDiagnosticsPayload) or isinstance(message, DiagnosticsPayload):
         kwargs = dict(rssi=message.rssi, enabled=message.flags.state)
-        if isinstance(message, InitialStatPayload):
+        if isinstance(message, InitialDiagnosticsPayload):
             kwargs['fw_version'] = message.fw_version
         relay_state.update(**kwargs)
 
@@ -157,8 +157,8 @@ def markup(ctx: Optional[bot.Context]) -> Optional[ReplyKeyboardMarkup]:
 
 
 if __name__ == '__main__':
-    mqtt_relay = MQTTRelay(devices=MQTTRelayDevice(id=config['mqtt']['home_id'],
-                                                   secret=config['mqtt']['home_secret']))
+    mqtt_relay = MqttRelay(devices=MqttEspDevice(id=config['mqtt']['home_id'],
+                                                 secret=config['mqtt']['home_secret']))
     mqtt_relay.set_message_callback(on_mqtt_message)
     mqtt_relay.configure_tls()
     mqtt_relay.connect_and_loop(loop_forever=False)