Add state_class support to Rest (#58026)

tests/components/rest/test_binary_sensor.py

@@ -10,6 +10,7 @@ import respx
 from homeassistant import config as hass_config
 import homeassistant.components.binary_sensor as binary_sensor
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     CONTENT_TYPE_JSON,
     SERVICE_RELOAD,
@@ -164,6 +165,7 @@ async def test_setup_get(hass):
                 "username": "my username",
                 "password": "my password",
                 "headers": {"Accept": CONTENT_TYPE_JSON},
+                "device_class": binary_sensor.DEVICE_CLASS_PLUG,
             }
         },
     )
@@ -171,6 +173,10 @@ async def test_setup_get(hass):
     await hass.async_block_till_done()
     assert len(hass.states.async_all("binary_sensor")) == 1
 
+    state = hass.states.get("binary_sensor.foo")
+    assert state.state == STATE_OFF
+    assert state.attributes[ATTR_DEVICE_CLASS] == binary_sensor.DEVICE_CLASS_PLUG
+
 
 @respx.mock
 async def test_setup_get_digest_auth(hass):

tests/components/rest/test_sensor.py

@@ -10,12 +10,15 @@ from homeassistant import config as hass_config
 from homeassistant.components.homeassistant import SERVICE_UPDATE_ENTITY
 import homeassistant.components.sensor as sensor
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     ATTR_UNIT_OF_MEASUREMENT,
     CONTENT_TYPE_JSON,
     DATA_MEGABYTES,
+    DEVICE_CLASS_TEMPERATURE,
     SERVICE_RELOAD,
     STATE_UNKNOWN,
+    TEMP_CELSIUS,
 )
 from homeassistant.setup import async_setup_component
 
@@ -177,13 +180,15 @@ async def test_setup_get(hass):
                 "method": "GET",
                 "value_template": "{{ value_json.key }}",
                 "name": "foo",
-                "unit_of_measurement": DATA_MEGABYTES,
+                "unit_of_measurement": TEMP_CELSIUS,
                 "verify_ssl": "true",
                 "timeout": 30,
                 "authentication": "basic",
                 "username": "my username",
                 "password": "my password",
                 "headers": {"Accept": CONTENT_TYPE_JSON},
+                "device_class": DEVICE_CLASS_TEMPERATURE,
+                "state_class": sensor.STATE_CLASS_MEASUREMENT,
             }
         },
     )
@@ -200,7 +205,11 @@ async def test_setup_get(hass):
         blocking=True,
     )
     await hass.async_block_till_done()
-    assert hass.states.get("sensor.foo").state == ""
+    state = hass.states.get("sensor.foo")
+    assert state.state == ""
+    assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == TEMP_CELSIUS
+    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_TEMPERATURE
+    assert state.attributes[sensor.ATTR_STATE_CLASS] == sensor.STATE_CLASS_MEASUREMENT
 
 
 @respx.mock

tests/components/rest/test_switch.py

@@ -6,7 +6,7 @@ import aiohttp
 
 from homeassistant.components.rest import DOMAIN
 import homeassistant.components.rest.switch as rest
-from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
+from homeassistant.components.switch import DEVICE_CLASS_SWITCH, DOMAIN as SWITCH_DOMAIN
 from homeassistant.const import (
     CONF_HEADERS,
     CONF_NAME,
@@ -23,6 +23,7 @@ from tests.common import assert_setup_component
 """Tests for setting up the REST switch platform."""
 
 NAME = "foo"
+DEVICE_CLASS = DEVICE_CLASS_SWITCH
 METHOD = "post"
 RESOURCE = "http://localhost/"
 STATE_RESOURCE = RESOURCE
@@ -158,6 +159,7 @@ def _setup_test_switch(hass):
     body_off = Template("off", hass)
     switch = rest.RestSwitch(
         NAME,
+        DEVICE_CLASS,
         RESOURCE,
         STATE_RESOURCE,
         METHOD,
@@ -180,6 +182,12 @@ def test_name(hass):
     assert switch.name == NAME
 
 
+def test_device_class(hass):
+    """Test the name."""
+    switch, body_on, body_off = _setup_test_switch(hass)
+    assert switch.device_class == DEVICE_CLASS
+
+
 def test_is_on_before_update(hass):
     """Test is_on in initial state."""
     switch, body_on, body_off = _setup_test_switch(hass)