1 files changed, 10 insertions, 9 deletions
diff --git a/src/relay_mqtt_http_proxy.py b/src/relay_mqtt_http_proxy.py
index 51a4e21..098facc 100755
--- a/src/relay_mqtt_http_proxy.py
+++ b/src/relay_mqtt_http_proxy.py
@@ -1,20 +1,21 @@
 #!/usr/bin/env python3
 from home import http
 from home.config import config
-from home.mqtt import MQTTRelay, MQTTRelayDevice, MQTTRelayState
-from home.mqtt.payload import MQTTPayload
-from home.mqtt.payload.relay import InitialStatPayload, StatPayload
+from home.mqtt import MqttRelay, MqttRelayState
+from home.mqtt.esp import MqttEspDevice
+from home.mqtt.payload import MqttPayload
+from home.mqtt.payload.relay import InitialDiagnosticsPayload, DiagnosticsPayload
 from typing import Optional
 
-mqtt_relay: Optional[MQTTRelay] = None
-relay_states: dict[str, MQTTRelayState] = {}
+mqtt_relay: Optional[MqttRelay] = None
+relay_states: dict[str, MqttRelayState] = {}
 
 
-def on_mqtt_message(device_id, message: MQTTPayload):
-    if isinstance(message, InitialStatPayload) or isinstance(message, StatPayload):
+def on_mqtt_message(device_id, message: MqttPayload):
+    if isinstance(message, InitialDiagnosticsPayload) or isinstance(message, DiagnosticsPayload):
         kwargs = dict(rssi=message.rssi, enabled=message.flags.state)
         if device_id not in relay_states:
-            relay_states[device_id] = MQTTRelayState()
+            relay_states[device_id] = MqttRelayState()
         relay_states[device_id].update(**kwargs)
 
 
@@ -54,7 +55,7 @@ class RelayMqttHttpProxy(http.HTTPServer):
 if __name__ == '__main__':
     config.load('relay_mqtt_http_proxy')
 
-    mqtt_relay = MQTTRelay(devices=[MQTTRelayDevice(id=device_id) for device_id in config.get('relay.devices')])
+    mqtt_relay = MqttRelay(devices=[MqttEspDevice(id=device_id) for device_id in config.get('relay.devices')])
     mqtt_relay.configure_tls()
     mqtt_relay.set_message_callback(on_mqtt_message)
     mqtt_relay.connect_and_loop(loop_forever=False)