diff --git a/src/driver/drv_main.c b/src/driver/drv_main.c
index 6d428a945..80c786bcd 100644
--- a/src/driver/drv_main.c
+++ b/src/driver/drv_main.c
@@ -24,6 +24,12 @@ const char *counter_mqttNames[OBK_NUM_COUNTERS] = {
     "consumption_stats"
 };
 
+const char *counter_devClasses[OBK_NUM_COUNTERS] = {
+    "energy",
+    "energy",
+    "None"
+};
+
 typedef struct driver_s {
 	const char *name;
 	void (*initFunc)();
diff --git a/src/driver/drv_public.h b/src/driver/drv_public.h
index afe85c19b..48e6eac58 100644
--- a/src/driver/drv_public.h
+++ b/src/driver/drv_public.h
@@ -22,6 +22,7 @@ enum {
 // MQTT names of sensors (voltage, current, power)
 extern const char *sensor_mqttNames[];
 extern const char *counter_mqttNames[];
+extern const char *counter_devClasses[];
 
 void DRV_Generic_Init();
 void DRV_AppendInformationToHTTPIndexPage(http_request_t *request);
diff --git a/src/driver/drv_tuyaMCU.c b/src/driver/drv_tuyaMCU.c
index 47ca4ef3f..175551db2 100644
--- a/src/driver/drv_tuyaMCU.c
+++ b/src/driver/drv_tuyaMCU.c
@@ -1017,11 +1017,13 @@ void TuyaMCU_RunFrame() {
     /* Command controll */
     if (heartbeat_timer == 0)
     {
+        /* Generate heartbeat to keep communication alove */
         TuyaMCU_SendCommandWithData(TUYA_CMD_HEARTBEAT, NULL, 0);
         heartbeat_timer = 3;
         heartbeat_counter++;
         if (heartbeat_counter>=4) 
         {
+            /* unanswerred heartbeats -> lost communication */
             heartbeat_valid = false;
             product_information_valid = false;
             working_mode_valid = false;
@@ -1029,6 +1031,7 @@ void TuyaMCU_RunFrame() {
             state_updated = false;
         }
     } else {
+        /* Heartbeat timer - sent every 3 seconds */
         if (heartbeat_timer>0)
         {
             heartbeat_timer--;
@@ -1037,25 +1040,32 @@ void TuyaMCU_RunFrame() {
         }
         if (heartbeat_valid == true)
         {
+            /* Connection Active */
             if (product_information_valid == false)
             {
+                /* Request production information */
                 TuyaMCU_SendCommandWithData(TUYA_CMD_QUERY_PRODUCT, NULL, 0);
             } 
             else if (working_mode_valid == false)
             {
+                /* Request working mode */
                 TuyaMCU_SendCommandWithData(TUYA_CMD_MCU_CONF, NULL, 0);
             }
             else if ((wifi_state_valid == false) && (self_processing_mode == false))
             {
+                /* Reset wifi state -> Aquirring network connection */ 
                 Tuya_SetWifiState(0);
-                //TuyaMCU_SendCommandWithData(TUYA_CMD_WIFI_STATE, NULL, 0);
+                TuyaMCU_SendCommandWithData(TUYA_CMD_WIFI_STATE, NULL, 0);
             }
             else if (state_updated == false)
             {
+                /* Request first state of all DP - this should list all existing DP */
                 TuyaMCU_SendCommandWithData(TUYA_CMD_QUERY_STATE, NULL, 0);
             }
             else 
             {
+                /* Monitor WIFI and MQTT connection and apply Wifi state 
+                 * State is updated when change is detected or after timeout */
                 if ((Main_HasWiFiConnected()!=0) && (Main_HasMQTTConnected()!=0))
                 {
                     if ((wifi_state == false) || (wifi_state_timer == 0))
@@ -1070,9 +1080,11 @@ void TuyaMCU_RunFrame() {
                         wifi_state = false;
                     }
                 }
+                /* wifi state timer */
                 wifi_state_timer++;
                 if (wifi_state_timer >= 60)
                 {
+                    /* timeout after ~1 minute */
                     wifi_state_timer = 0;
                 }
             }
diff --git a/src/httpserver/hass.c b/src/httpserver/hass.c
index b655b457c..9ee80cdef 100644
--- a/src/httpserver/hass.c
+++ b/src/httpserver/hass.c
@@ -243,7 +243,7 @@ HassDeviceInfo* hass_init_light_device_info(ENTITY_TYPE type, int index) {
         }
         if ((index >= OBK_CONSUMPTION_TOTAL) && (index <= OBK_CONSUMPTION_STATS))
         {
-            cJSON_AddStringToObject(info->root, "dev_cla", counter_mqttNames[index - OBK_CONSUMPTION_TOTAL]);  //device_class=consumption
+            cJSON_AddStringToObject(info->root, "dev_cla", counter_devClasses[index - OBK_CONSUMPTION_TOTAL]);  //device_class=consumption
 
             sprintf(g_hassBuffer, "%s/%s/get", clientId, counter_mqttNames[index - OBK_CONSUMPTION_TOTAL]);
             cJSON_AddStringToObject(info->root, STATE_TOPIC_KEY, g_hassBuffer);
diff --git a/src/httpserver/http_fns.c b/src/httpserver/http_fns.c
index dcc82f89b..b382fb9a0 100644
--- a/src/httpserver/http_fns.c
+++ b/src/httpserver/http_fns.c
@@ -1322,7 +1322,7 @@ int http_fn_ha_discovery(http_request_t* request) {
 
 #ifndef OBK_DISABLE_ALL_DRIVERS
 	if (measuringPower == true) {
-		for (i = 0;i < OBK_NUM_MEASUREMENTS;i++)
+		for (i = 0;i < OBK_NUM_SENSOR_COUNT;i++)
 		{
 			HassDeviceInfo* dev_info = hass_init_light_device_info(ENTITY_SENSOR, i);
 			MQTT_QueuePublish(topic, dev_info->channel, hass_build_discovery_json(dev_info), OBK_PUBLISH_FLAG_RETAIN);