diff --git a/library.json b/library.json
index b6f29eb..e79e522 100644
--- a/library.json
+++ b/library.json
@@ -1,7 +1,7 @@
 {
   "$schema": "https://raw.githubusercontent.com/platformio/platformio-core/develop/platformio/assets/schema/library.json",
   "name": "ha-mqtt",
-  "version": "1.5.1",
+  "version": "1.6.0",
   "description": "Home Assistant classes for integration with MQTT auto discovery",
   "repository": {
     "type": "git",
diff --git a/src/esp.h b/src/esp.h
index 5cfb8a7..4dd511b 100644
--- a/src/esp.h
+++ b/src/esp.h
@@ -60,6 +60,15 @@ namespace HaESP {
         return builder;
     }
 
+    Builder<Button>& restartButton() {
+        return Builder<Button>::instance(new Button{"Restart", "restart",
+            [](const char* msg) {
+                if (strcmp("PRESS", msg) == 0) ESP.restart();
+            }
+        })
+        .withIcon("mdi:restart");
+    }
+
     void enableSensors() {
         if (activeSensors.heapStats) tHeap.enable();
         if (activeSensors.restartInfo) tRestartInfo.enable();
diff --git a/src/ha.h b/src/ha.h
index 3d627ea..56d182e 100644
--- a/src/ha.h
+++ b/src/ha.h
@@ -77,6 +77,7 @@ namespace Ha {
         const char* name = nullptr;
         const char* entityCategory = nullptr;
         const char* deviceClass = nullptr;
+        const char* icon = nullptr;
         DeviceConfig* mainDevice = nullptr;
         inline static List<Component> components;
         bool multiValueComponent = false;
@@ -104,6 +105,7 @@ namespace Ha {
             if (mainDevice) mainDevice->buildConfig(jsonDoc);
             if (entityCategory) jsonDoc["entity_category"] = entityCategory;
             if (deviceClass) jsonDoc["device_class"] = deviceClass;
+            if (icon) jsonDoc["icon"] = icon;
             jsonDoc["name"] = name;
            
             buildUniqueId(jsonDoc);
@@ -159,6 +161,7 @@ namespace Ha {
     struct State : Config {
         char stateTopic[TOPIC_SIZE] = {};
         const char* jsonAttributesTemplate = nullptr;
+        const char* valueTemplate = nullptr;
 
         State(Component* cmp) : cmp(cmp) {}
 
@@ -178,6 +181,7 @@ namespace Ha {
                     jsonDoc["json_attributes_template"] = jsonAttributesTemplate;
                     jsonDoc["json_attributes_topic"] = (const char*)stateTopic;
                 }
+                if (valueTemplate) jsonDoc["value_template"] = valueTemplate;
             }
         }
     private:
@@ -243,6 +247,7 @@ namespace Ha {
         unsigned int min = 1;
         unsigned int max = 100;
         unsigned int step = 1;
+        const char* unitMeasure = nullptr;
 
         Number(const char* name, const char* id, onMessage f) : Component(id, name, "number"), StatefulCommand(this, f) {}
 
@@ -256,6 +261,7 @@ namespace Ha {
             jsonDoc["min"] = min;
             jsonDoc["max"] = max;
             jsonDoc["step"] = step;
+            if (unitMeasure) jsonDoc["unit_of_measurement"] = unitMeasure;
         }
     };
 
@@ -282,7 +288,6 @@ namespace Ha {
 
     struct Sensor : Component, State {
         const char* unitMeasure = nullptr;
-        const char* valueTemplate = nullptr;
         unsigned int precision = 2;
         inline static unordered_map<string, Sensor*> mapSensors;
 
@@ -295,7 +300,6 @@ namespace Ha {
             Component::buildConfig(jsonDoc);
             State::buildConfig(jsonDoc);
             if (unitMeasure) jsonDoc["unit_of_measurement"] = unitMeasure;
-            if (valueTemplate) jsonDoc["value_template"] = valueTemplate;
             if (isNumericSensor()) jsonDoc["suggested_display_precision"] = precision;
         }
 
@@ -388,6 +392,11 @@ namespace Ha {
             return *this;
         }
 
+        Builder& withIcon(const char* value) {
+            cmp->icon = value;
+            return *this;
+        }
+
         Builder& withValueTemplate(const char* value) {
             cmp->valueTemplate = value;
             return *this;