diff --git a/GridboxConnectorAddon-dev/GridboxConnector/ha_viessmann_gridbox_connector.py b/GridboxConnectorAddon-dev/GridboxConnector/ha_viessmann_gridbox_connector.py
index 2baf7d1..123b176 100644
--- a/GridboxConnectorAddon-dev/GridboxConnector/ha_viessmann_gridbox_connector.py
+++ b/GridboxConnectorAddon-dev/GridboxConnector/ha_viessmann_gridbox_connector.py
@@ -81,7 +81,9 @@ def __init__(self, mqtt_settings):
         # Battery sum
         self.battery_sum = HAViessmannBattery(mqtt_settings, self.device_info, "sum", "")
 
+        # Heater
         self.heater_sensor = HAViessmannHeater(mqtt_settings, self.device_info, "", "")
+        
         # EV
         self.ev_sum = HAViessmannEVChargingStation(mqtt_settings, self.device_info, "sum", "")
 
diff --git a/GridboxConnectorAddon-dev/GridboxConnector/main_test.py b/GridboxConnectorAddon-dev/GridboxConnector/main_test.py
index 373e6cc..74b78f8 100644
--- a/GridboxConnectorAddon-dev/GridboxConnector/main_test.py
+++ b/GridboxConnectorAddon-dev/GridboxConnector/main_test.py
@@ -7,9 +7,6 @@
 from paho.mqtt.client import MQTT_ERR_SUCCESS
 from ha_mqtt_discoverable.sensors import Sensor
 
-mock_data = "{'batteries': [{'applianceID': '97af25d0-3791-48cc-857c-14aaac749267', 'capacity': 10000, 'nominalCapacity': 10000, 'power': -853, 'remainingCharge': 7700, 'stateOfCharge': 0.77}], 'battery': {'capacity': 10000, 'nominalCapacity': 10000, 'power': -853, 'remainingCharge': 7700, 'stateOfCharge': 0.77}, 'consumption': 600, 'directConsumption': 600, 'directConsumptionEV': 0, 'directConsumptionHeatPump': 0, 'directConsumptionHeater': 0, 'directConsumptionHousehold': 600, 'directConsumptionRate': 0.3968253968253968, 'grid': -59, 'gridMeterReadingNegative': 14081760000, 'gridMeterReadingPositive': 7393320000, 'measuredAt': '2024-05-08T09:42:18Z', 'photovoltaic': 1512, 'production': 1512, 'selfConsumption': 1453, 'selfConsumptionRate': 0.9609788359788359, 'selfSufficiencyRate': 1, 'selfSupply': 600, 'totalConsumption': 600}"
-
-
 class TestGridboxConnectorMethods(unittest.TestCase):
 
     @patch('paho.mqtt.client.Client')
@@ -19,7 +16,7 @@ class TestGridboxConnectorMethods(unittest.TestCase):
     def test_main_batteries(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_mqtt_client):
         # Load mock data from JSON file
         with open('tests/mock_data/mock_data_with_batteries.json') as f:
-            mock_data = json.load(f)
+            mock_data = [json.load(f)]
         mock_mqtt_client.return_value.connect.return_value = MQTT_ERR_SUCCESS
         mock_retrieve_live_data.return_value = mock_data
         # Create an instance of the class
@@ -50,11 +47,11 @@ def test_main_batteries(self, mock_retrieve_live_data, mock_init, mock_init_auth
              patch.object(viessmann_gridbox_connector.self_sufficiency_rate_sensor, 'set_state') as mock_self_sufficiency_rate_sensor:
             viessmann_gridbox_connector.update_sensors(result[0])
             mock_self_supply_sensor.assert_called_once_with(600.0)
-            mock_self_consumtion_rate_sensor.assert_called_once_with(96.09788359788359)
+            mock_self_consumtion_rate_sensor.assert_called_once_with(96.1)
             mock_consumption_household_sensor.assert_called_once_with(600)
             mock_direct_consumption_household_sensor.assert_called_once_with(600.0)
             mock_direct_consumption_heatpump_sensor.assert_called_once_with(0.0)
-            mock_direct_consumption_rate_sensor.assert_called_once_with(39.682539682539684)
+            mock_direct_consumption_rate_sensor.assert_called_once_with(39.68)
             mock_self_sufficiency_rate_sensor.assert_called_once_with(100.0)
 
     @patch('paho.mqtt.client.Client')
@@ -64,7 +61,7 @@ def test_main_batteries(self, mock_retrieve_live_data, mock_init, mock_init_auth
     def test_main_ev(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_mqtt_client):
         # Load mock data from JSON file
         with open('tests/mock_data/mock_data_with_ev.json') as f:
-            mock_data = json.load(f)
+            mock_data = [json.load(f)]
         mock_mqtt_client.return_value.connect.return_value = MQTT_ERR_SUCCESS
         mock_retrieve_live_data.return_value = mock_data
         # Create an instance of the class
@@ -86,23 +83,42 @@ def test_main_ev(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_m
         viessmann_gridbox_connector = HAViessmannGridboxConnector(
             mqtt_settings)
 
-        with patch.object(viessmann_gridbox_connector.self_supply_sensor, 'set_state') as mock_self_supply_sensor, \
-             patch.object(viessmann_gridbox_connector.self_consumtion_rate_sensor, 'set_state') as mock_self_consumption_rate_sensor, \
-             patch.object(viessmann_gridbox_connector.consumption_household_sensor, 'set_state') as mock_consumption_household_sensor, \
-             patch.object(viessmann_gridbox_connector.direct_consumption_household_sensor, 'set_state') as mock_direct_consumption_household_sensor, \
-             patch.object(viessmann_gridbox_connector.direct_consumption_heatpump_sensor, 'set_state') as mock_direct_consumption_heatpump_sensor, \
-             patch.object(viessmann_gridbox_connector.direct_consumption_rate_sensor, 'set_state') as mock_direct_consumption_rate_sensor, \
-             patch.object(viessmann_gridbox_connector.self_sufficiency_rate_sensor, 'set_state') as mock_self_sufficiency_rate_sensor, \
-             patch.object(viessmann_gridbox_connector.ev_sum_sensor, 'set_state') as mock_ev_sum_sensor:
+        with patch.object(viessmann_gridbox_connector.ev_sum, 'set_states') as mock_ev_sum_sensor:
             viessmann_gridbox_connector.update_sensors(result[0])
-            mock_self_supply_sensor.assert_called_once_with(8854)
-            mock_self_consumption_rate_sensor.assert_called_once_with(96.09788359788359)
-            mock_consumption_household_sensor.assert_called_once_with(600)
-            mock_direct_consumption_household_sensor.assert_called_once_with(600.0)
-            mock_direct_consumption_heatpump_sensor.assert_called_once_with(0.0)
-            mock_direct_consumption_rate_sensor.assert_called_once_with(0.8414867048626414)
-            mock_self_sufficiency_rate_sensor.assert_called_once_with(100.0)
             mock_ev_sum_sensor.assert_called_once_with(7930.4,0,0,0,0,0)
 
+    @patch('paho.mqtt.client.Client')
+    @patch.object(GridboxConnector, 'init_auth', return_value=None)
+    @patch.object(GridboxConnector, '__init__', return_value=None)
+    @patch.object(GridboxConnector, 'retrieve_live_data')
+    def test_main_heater(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_mqtt_client):
+        # Load mock data from JSON file
+        with open('tests/mock_data/mock_data_with_heater.json') as f:
+            mock_data = [json.load(f)]
+        mock_mqtt_client.return_value.connect.return_value = MQTT_ERR_SUCCESS
+        mock_retrieve_live_data.return_value = mock_data
+        # Create an instance of the class
+        gridbox_connector = GridboxConnector(None)
+
+        # Call the function
+        result = gridbox_connector.retrieve_live_data()
+
+        # Assert the function was called once
+        mock_retrieve_live_data.assert_called_once()
+
+        # Assert the function returned the mock value
+        self.assertEqual(result, mock_data)
+        mqtt_server = "mqtt_server"
+        mqtt_user = "mqtt_user"
+        mqtt_pw = "mqtt_pw"
+        mqtt_settings = Settings.MQTT(
+            host=mqtt_server, username=mqtt_user, password=mqtt_pw)
+        viessmann_gridbox_connector = HAViessmannGridboxConnector(
+            mqtt_settings)
+
+        with patch.object(viessmann_gridbox_connector.heater_sensor, 'set_states') as mock_heater_sensor:
+            viessmann_gridbox_connector.update_sensors(result[0])
+            mock_heater_sensor.assert_called_once_with(3676,70.9)
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/GridboxConnectorAddon-edge/GridboxConnector/ha_viessmann_gridbox_connector.py b/GridboxConnectorAddon-edge/GridboxConnector/ha_viessmann_gridbox_connector.py
index 2baf7d1..123b176 100644
--- a/GridboxConnectorAddon-edge/GridboxConnector/ha_viessmann_gridbox_connector.py
+++ b/GridboxConnectorAddon-edge/GridboxConnector/ha_viessmann_gridbox_connector.py
@@ -81,7 +81,9 @@ def __init__(self, mqtt_settings):
         # Battery sum
         self.battery_sum = HAViessmannBattery(mqtt_settings, self.device_info, "sum", "")
 
+        # Heater
         self.heater_sensor = HAViessmannHeater(mqtt_settings, self.device_info, "", "")
+        
         # EV
         self.ev_sum = HAViessmannEVChargingStation(mqtt_settings, self.device_info, "sum", "")
 
diff --git a/GridboxConnectorAddon-edge/GridboxConnector/main_test.py b/GridboxConnectorAddon-edge/GridboxConnector/main_test.py
index 373e6cc..7704f3b 100644
--- a/GridboxConnectorAddon-edge/GridboxConnector/main_test.py
+++ b/GridboxConnectorAddon-edge/GridboxConnector/main_test.py
@@ -7,9 +7,6 @@
 from paho.mqtt.client import MQTT_ERR_SUCCESS
 from ha_mqtt_discoverable.sensors import Sensor
 
-mock_data = "{'batteries': [{'applianceID': '97af25d0-3791-48cc-857c-14aaac749267', 'capacity': 10000, 'nominalCapacity': 10000, 'power': -853, 'remainingCharge': 7700, 'stateOfCharge': 0.77}], 'battery': {'capacity': 10000, 'nominalCapacity': 10000, 'power': -853, 'remainingCharge': 7700, 'stateOfCharge': 0.77}, 'consumption': 600, 'directConsumption': 600, 'directConsumptionEV': 0, 'directConsumptionHeatPump': 0, 'directConsumptionHeater': 0, 'directConsumptionHousehold': 600, 'directConsumptionRate': 0.3968253968253968, 'grid': -59, 'gridMeterReadingNegative': 14081760000, 'gridMeterReadingPositive': 7393320000, 'measuredAt': '2024-05-08T09:42:18Z', 'photovoltaic': 1512, 'production': 1512, 'selfConsumption': 1453, 'selfConsumptionRate': 0.9609788359788359, 'selfSufficiencyRate': 1, 'selfSupply': 600, 'totalConsumption': 600}"
-
-
 class TestGridboxConnectorMethods(unittest.TestCase):
 
     @patch('paho.mqtt.client.Client')
@@ -19,7 +16,7 @@ class TestGridboxConnectorMethods(unittest.TestCase):
     def test_main_batteries(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_mqtt_client):
         # Load mock data from JSON file
         with open('tests/mock_data/mock_data_with_batteries.json') as f:
-            mock_data = json.load(f)
+            mock_data = [json.load(f)]
         mock_mqtt_client.return_value.connect.return_value = MQTT_ERR_SUCCESS
         mock_retrieve_live_data.return_value = mock_data
         # Create an instance of the class
@@ -50,11 +47,11 @@ def test_main_batteries(self, mock_retrieve_live_data, mock_init, mock_init_auth
              patch.object(viessmann_gridbox_connector.self_sufficiency_rate_sensor, 'set_state') as mock_self_sufficiency_rate_sensor:
             viessmann_gridbox_connector.update_sensors(result[0])
             mock_self_supply_sensor.assert_called_once_with(600.0)
-            mock_self_consumtion_rate_sensor.assert_called_once_with(96.09788359788359)
+            mock_self_consumtion_rate_sensor.assert_called_once_with(96.1)
             mock_consumption_household_sensor.assert_called_once_with(600)
             mock_direct_consumption_household_sensor.assert_called_once_with(600.0)
             mock_direct_consumption_heatpump_sensor.assert_called_once_with(0.0)
-            mock_direct_consumption_rate_sensor.assert_called_once_with(39.682539682539684)
+            mock_direct_consumption_rate_sensor.assert_called_once_with(39.68)
             mock_self_sufficiency_rate_sensor.assert_called_once_with(100.0)
 
     @patch('paho.mqtt.client.Client')
@@ -64,7 +61,7 @@ def test_main_batteries(self, mock_retrieve_live_data, mock_init, mock_init_auth
     def test_main_ev(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_mqtt_client):
         # Load mock data from JSON file
         with open('tests/mock_data/mock_data_with_ev.json') as f:
-            mock_data = json.load(f)
+            mock_data = [json.load(f)]
         mock_mqtt_client.return_value.connect.return_value = MQTT_ERR_SUCCESS
         mock_retrieve_live_data.return_value = mock_data
         # Create an instance of the class
@@ -86,23 +83,42 @@ def test_main_ev(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_m
         viessmann_gridbox_connector = HAViessmannGridboxConnector(
             mqtt_settings)
 
-        with patch.object(viessmann_gridbox_connector.self_supply_sensor, 'set_state') as mock_self_supply_sensor, \
-             patch.object(viessmann_gridbox_connector.self_consumtion_rate_sensor, 'set_state') as mock_self_consumption_rate_sensor, \
-             patch.object(viessmann_gridbox_connector.consumption_household_sensor, 'set_state') as mock_consumption_household_sensor, \
-             patch.object(viessmann_gridbox_connector.direct_consumption_household_sensor, 'set_state') as mock_direct_consumption_household_sensor, \
-             patch.object(viessmann_gridbox_connector.direct_consumption_heatpump_sensor, 'set_state') as mock_direct_consumption_heatpump_sensor, \
-             patch.object(viessmann_gridbox_connector.direct_consumption_rate_sensor, 'set_state') as mock_direct_consumption_rate_sensor, \
-             patch.object(viessmann_gridbox_connector.self_sufficiency_rate_sensor, 'set_state') as mock_self_sufficiency_rate_sensor, \
-             patch.object(viessmann_gridbox_connector.ev_sum_sensor, 'set_state') as mock_ev_sum_sensor:
+        with patch.object(viessmann_gridbox_connector.ev_sum_sensor, 'set_state') as mock_ev_sum_sensor:
             viessmann_gridbox_connector.update_sensors(result[0])
-            mock_self_supply_sensor.assert_called_once_with(8854)
-            mock_self_consumption_rate_sensor.assert_called_once_with(96.09788359788359)
-            mock_consumption_household_sensor.assert_called_once_with(600)
-            mock_direct_consumption_household_sensor.assert_called_once_with(600.0)
-            mock_direct_consumption_heatpump_sensor.assert_called_once_with(0.0)
-            mock_direct_consumption_rate_sensor.assert_called_once_with(0.8414867048626414)
-            mock_self_sufficiency_rate_sensor.assert_called_once_with(100.0)
             mock_ev_sum_sensor.assert_called_once_with(7930.4,0,0,0,0,0)
 
+    @patch('paho.mqtt.client.Client')
+    @patch.object(GridboxConnector, 'init_auth', return_value=None)
+    @patch.object(GridboxConnector, '__init__', return_value=None)
+    @patch.object(GridboxConnector, 'retrieve_live_data')
+    def test_main_heater(self, mock_retrieve_live_data, mock_init, mock_init_auth,mock_mqtt_client):
+        # Load mock data from JSON file
+        with open('tests/mock_data/mock_data_with_heater.json') as f:
+            mock_data = [json.load(f)]
+        mock_mqtt_client.return_value.connect.return_value = MQTT_ERR_SUCCESS
+        mock_retrieve_live_data.return_value = mock_data
+        # Create an instance of the class
+        gridbox_connector = GridboxConnector(None)
+
+        # Call the function
+        result = gridbox_connector.retrieve_live_data()
+
+        # Assert the function was called once
+        mock_retrieve_live_data.assert_called_once()
+
+        # Assert the function returned the mock value
+        self.assertEqual(result, mock_data)
+        mqtt_server = "mqtt_server"
+        mqtt_user = "mqtt_user"
+        mqtt_pw = "mqtt_pw"
+        mqtt_settings = Settings.MQTT(
+            host=mqtt_server, username=mqtt_user, password=mqtt_pw)
+        viessmann_gridbox_connector = HAViessmannGridboxConnector(
+            mqtt_settings)
+
+        with patch.object(viessmann_gridbox_connector.heater_sensor, 'set_state') as mock_heater_sensor:
+            viessmann_gridbox_connector.update_sensors(result[0])
+            mock_heater_sensor.assert_called_once_with(3676,70.9)
+
 if __name__ == '__main__':
     unittest.main()