diff --git a/power_mqtt1.py b/power_mqtt1.py
index 7d97679..3f2bb44 100755
--- a/power_mqtt1.py
+++ b/power_mqtt1.py
@@ -153,9 +153,12 @@ class PowerCTL:
     def mqtt_send_event(self, relay: ButtonInput, ev):
         self.mqtc.publish(f"{self.PREFIX}/{relay.name}/input", ev, qos=0, retain=False)
 
-    @staticmethod
-    def on_connect(_mosq, _obj, _foo, _bar):
-        print("Connected")
+    def on_connect(self, _mosq, _obj, connect_flags, _rc):
+        print("Connected", connect_flags)
+        if not connect_flags.get("session present", 0):
+            print("Re-Subscribe")
+            for topic in self.relays:
+                self.mqtc.subscribe(topic, 0)
 
     @staticmethod
     def on_subscribe(_mosq, _obj, _mid, _granted_qos):
diff --git a/test.py b/test.py
index 922d442..38a792a 100644
--- a/test.py
+++ b/test.py
@@ -9,21 +9,26 @@ import power_mqtt1
 
 class SMBus:
     def __init__(self, _bus):
-        pass
+        self.storage = {}
 
     def read_byte(self, _i2c_addr):
         time.sleep(5)
-        return random.randbytes(1)[0]
+        # return random.randbytes(1)[0]
+        return self.storage.get(_i2c_addr, -1)
 
     def write_byte(self, i2c_addr, data):
-        pass
+        self.storage[i2c_addr] = data
 
 
 class MqttClient(mqtt.Client):
     def connect(self, host, *args, **kwargs):
         return super().connect("127.0.0.1", *args, **kwargs)
+        # return super().connect(host, *args, **kwargs)
+
+    def publish(self, topic, payload=None, qos=0, retain=False, properties=None, ):
+        pass
 
 
-with patch('smbus.SMBus', new=SMBus) as SMBus_mock, \
-        patch("paho.mqtt.client.Client", new=MqttClient) as connect_mock:
-    power_mqtt1.main()
+with patch('smbus.SMBus', new=SMBus) as SMBus_mock:
+    with patch("paho.mqtt.client.Client", new=MqttClient) as connect_mock:
+        power_mqtt1.main()