Add device class support to text sensor (esphome#6202)

Co-authored-by: Jesse Hills <[email protected]>

esphome/components/api/api.proto

@@ -600,6 +600,7 @@ message ListEntitiesTextSensorResponse {
   string icon = 5;
   bool disabled_by_default = 6;
   EntityCategory entity_category = 7;
+  string device_class = 8;
 }
 message TextSensorStateResponse {
   option (id) = 27;

esphome/components/api/api_connection.cpp

@@ -543,6 +543,7 @@ bool APIConnection::send_text_sensor_info(text_sensor::TextSensor *text_sensor)
   msg.icon = text_sensor->get_icon();
   msg.disabled_by_default = text_sensor->is_disabled_by_default();
   msg.entity_category = static_cast<enums::EntityCategory>(text_sensor->get_entity_category());
+  msg.device_class = text_sensor->get_device_class();
   return this->send_list_entities_text_sensor_response(msg);
 }
 #endif

esphome/components/api/api_pb2.cpp

@@ -2721,6 +2721,10 @@ bool ListEntitiesTextSensorResponse::decode_length(uint32_t field_id, ProtoLengt
       this->icon = value.as_string();
       return true;
     }
+    case 8: {
+      this->device_class = value.as_string();
+      return true;
+    }
     default:
       return false;
   }
@@ -2743,6 +2747,7 @@ void ListEntitiesTextSensorResponse::encode(ProtoWriteBuffer buffer) const {
   buffer.encode_string(5, this->icon);
   buffer.encode_bool(6, this->disabled_by_default);
   buffer.encode_enum<enums::EntityCategory>(7, this->entity_category);
+  buffer.encode_string(8, this->device_class);
 }
 #ifdef HAS_PROTO_MESSAGE_DUMP
 void ListEntitiesTextSensorResponse::dump_to(std::string &out) const {
@@ -2776,6 +2781,10 @@ void ListEntitiesTextSensorResponse::dump_to(std::string &out) const {
   out.append("  entity_category: ");
   out.append(proto_enum_to_string<enums::EntityCategory>(this->entity_category));
   out.append("\n");
+
+  out.append("  device_class: ");
+  out.append("'").append(this->device_class).append("'");
+  out.append("\n");
   out.append("}");
 }
 #endif

esphome/components/api/api_pb2.h

@@ -713,6 +713,7 @@ class ListEntitiesTextSensorResponse : public ProtoMessage {
   std::string icon{};
   bool disabled_by_default{false};
   enums::EntityCategory entity_category{};
+  std::string device_class{};
   void encode(ProtoWriteBuffer buffer) const override;
 #ifdef HAS_PROTO_MESSAGE_DUMP
   void dump_to(std::string &out) const override;

esphome/components/mqtt/mqtt_text_sensor.cpp

@@ -1,6 +1,8 @@
 #include "mqtt_text_sensor.h"
 #include "esphome/core/log.h"
 
+#include "mqtt_const.h"
+
 #ifdef USE_MQTT
 #ifdef USE_TEXT_SENSOR
 
@@ -13,6 +15,8 @@ using namespace esphome::text_sensor;
 
 MQTTTextSensor::MQTTTextSensor(TextSensor *sensor) : sensor_(sensor) {}
 void MQTTTextSensor::send_discovery(JsonObject root, mqtt::SendDiscoveryConfig &config) {
+  if (!this->sensor_->get_device_class().empty())
+    root[MQTT_DEVICE_CLASS] = this->sensor_->get_device_class();
   config.command_topic = false;
 }
 void MQTTTextSensor::setup() {

esphome/components/text_sensor/__init__.py

@@ -3,6 +3,7 @@
 from esphome import automation
 from esphome.components import mqtt
 from esphome.const import (
+    CONF_DEVICE_CLASS,
     CONF_ENTITY_CATEGORY,
     CONF_FILTERS,
     CONF_ICON,
@@ -14,12 +15,21 @@
     CONF_STATE,
     CONF_FROM,
     CONF_TO,
+    DEVICE_CLASS_DATE,
+    DEVICE_CLASS_EMPTY,
+    DEVICE_CLASS_TIMESTAMP,
 )
 from esphome.core import CORE, coroutine_with_priority
 from esphome.cpp_generator import MockObjClass
 from esphome.cpp_helpers import setup_entity
 from esphome.util import Registry
 
+DEVICE_CLASSES = [
+    DEVICE_CLASS_DATE,
+    DEVICE_CLASS_EMPTY,
+    DEVICE_CLASS_TIMESTAMP,
+]
+
 
 IS_PLATFORM_COMPONENT = True
 
@@ -112,10 +122,13 @@ async def map_filter_to_code(config, filter_id):
     )
 
 
+validate_device_class = cv.one_of(*DEVICE_CLASSES, lower=True, space="_")
+
 TEXT_SENSOR_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMPONENT_SCHEMA).extend(
     {
         cv.OnlyWith(CONF_MQTT_ID, "mqtt"): cv.declare_id(mqtt.MQTTTextSensor),
         cv.GenerateID(): cv.declare_id(TextSensor),
+        cv.Optional(CONF_DEVICE_CLASS): validate_device_class,
         cv.Optional(CONF_FILTERS): validate_filters,
         cv.Optional(CONF_ON_VALUE): automation.validate_automation(
             {
@@ -140,12 +153,21 @@ def text_sensor_schema(
     *,
     icon: str = _UNDEF,
     entity_category: str = _UNDEF,
+    device_class: str = _UNDEF,
 ) -> cv.Schema:
     schema = TEXT_SENSOR_SCHEMA
     if class_ is not _UNDEF:
         schema = schema.extend({cv.GenerateID(): cv.declare_id(class_)})
     if icon is not _UNDEF:
         schema = schema.extend({cv.Optional(CONF_ICON, default=icon): cv.icon})
+    if device_class is not _UNDEF:
+        schema = schema.extend(
+            {
+                cv.Optional(
+                    CONF_DEVICE_CLASS, default=device_class
+                ): validate_device_class
+            }
+        )
     if entity_category is not _UNDEF:
         schema = schema.extend(
             {
@@ -164,6 +186,9 @@ async def build_filters(config):
 async def setup_text_sensor_core_(var, config):
     await setup_entity(var, config)
 
+    if CONF_DEVICE_CLASS in config:
+        cg.add(var.set_device_class(config[CONF_DEVICE_CLASS]))
+
     if config.get(CONF_FILTERS):  # must exist and not be empty
         filters = await build_filters(config[CONF_FILTERS])
         cg.add(var.set_filters(filters))

esphome/components/text_sensor/text_sensor.h

@@ -13,6 +13,9 @@ namespace text_sensor {
 #define LOG_TEXT_SENSOR(prefix, type, obj) \
   if ((obj) != nullptr) { \
     ESP_LOGCONFIG(TAG, "%s%s '%s'", prefix, LOG_STR_LITERAL(type), (obj)->get_name().c_str()); \
+    if (!(obj)->get_device_class().empty()) { \
+      ESP_LOGCONFIG(TAG, "%s  Device Class: '%s'", prefix, (obj)->get_device_class().c_str()); \
+    } \
     if (!(obj)->get_icon().empty()) { \
       ESP_LOGCONFIG(TAG, "%s  Icon: '%s'", prefix, (obj)->get_icon().c_str()); \
     } \
@@ -28,7 +31,7 @@ namespace text_sensor {
  public: \
   void set_##name##_text_sensor(text_sensor::TextSensor *text_sensor) { this->name##_text_sensor_ = text_sensor; }
 
-class TextSensor : public EntityBase {
+class TextSensor : public EntityBase, public EntityBase_DeviceClass {
  public:
   /// Getter-syntax for .state.
   std::string get_state() const;

tests/component_tests/text_sensor/test_text_sensor.py

@@ -0,0 +1,58 @@
+"""Tests for the text sensor component."""
+
+
+def test_text_sensor_is_setup(generate_main):
+    """
+    When the text is set in the yaml file, it should be registered in main
+    """
+    # Given
+
+    # When
+    main_cpp = generate_main("tests/component_tests/text_sensor/test_text_sensor.yaml")
+
+    # Then
+    assert "new template_::TemplateTextSensor();" in main_cpp
+    assert "App.register_text_sensor" in main_cpp
+
+
+def test_text_sensor_sets_mandatory_fields(generate_main):
+    """
+    When the mandatory fields are set in the yaml, they should be set in main
+    """
+    # Given
+
+    # When
+    main_cpp = generate_main("tests/component_tests/text_sensor/test_text_sensor.yaml")
+
+    # Then
+    assert 'ts_1->set_name("Template Text Sensor 1");' in main_cpp
+    assert 'ts_2->set_name("Template Text Sensor 2");' in main_cpp
+    assert 'ts_3->set_name("Template Text Sensor 3");' in main_cpp
+
+
+def test_text_sensor_config_value_internal_set(generate_main):
+    """
+    Test that the "internal" config value is correctly set
+    """
+    # Given
+
+    # When
+    main_cpp = generate_main("tests/component_tests/text_sensor/test_text_sensor.yaml")
+
+    # Then
+    assert "ts_2->set_internal(true);" in main_cpp
+    assert "ts_3->set_internal(false);" in main_cpp
+
+
+def test_text_sensor_device_class_set(generate_main):
+    """
+    When the device_class of text_sensor is set in the yaml file, it should be registered in main
+    """
+    # Given
+
+    # When
+    main_cpp = generate_main("tests/component_tests/text_sensor/test_text_sensor.yaml")
+
+    # Then
+    assert 'ts_2->set_device_class("timestamp");' in main_cpp
+    assert 'ts_3->set_device_class("date");' in main_cpp

tests/component_tests/text_sensor/test_text_sensor.yaml

@@ -0,0 +1,26 @@
+---
+esphome:
+  name: test
+  platform: ESP8266
+  board: d1_mini_lite
+
+text_sensor:
+  - platform: template
+    id: ts_1
+    name: "Template Text Sensor 1"
+    lambda: |-
+      return {"Hello World"};
+  - platform: template
+    id: ts_2
+    name: "Template Text Sensor 2"
+    lambda: |-
+      return {"2023-06-22T18:43:52+00:00"};
+    device_class: timestamp
+    internal: true
+  - platform: template
+    id: ts_3
+    name: "Template Text Sensor 3"
+    lambda: |-
+      return {"2023-06-22T18:43:52+00:00"};
+    device_class: date
+    internal: false

tests/test1.yaml

@@ -3923,6 +3923,10 @@ text_sensor:
   - platform: template
     name: Template Text Sensor
     id: ${textname}_text
+  - platform: template
+    name: Template Text Sensor Timestamp
+    id: ${textname}_text_timestamp
+    device_class: timestamp
   - platform: wifi_info
     scan_results:
       name: Scan Results