diff --git a/tests/components/wled/snapshots/test_button.ambr b/tests/components/wled/snapshots/test_button.ambr
new file mode 100644
index 00000000000..da487b49489
--- /dev/null
+++ b/tests/components/wled/snapshots/test_button.ambr
@@ -0,0 +1,75 @@
+# serializer version: 1
+# name: test_button_restart
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'device_class': 'restart',
+      'friendly_name': 'WLED RGB Light Restart',
+    }),
+    'context': <ANY>,
+    'entity_id': 'button.wled_rgb_light_restart',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'unknown',
+  })
+# ---
+# name: test_button_restart.1
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'button',
+    'entity_category': <EntityCategory.CONFIG: 'config'>,
+    'entity_id': 'button.wled_rgb_light_restart',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': <ButtonDeviceClass.RESTART: 'restart'>,
+    'original_icon': None,
+    'original_name': 'Restart',
+    'platform': 'wled',
+    'supported_features': 0,
+    'translation_key': None,
+    'unique_id': 'aabbccddeeff_restart',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_button_restart.2
+  DeviceRegistryEntrySnapshot({
+    'area_id': None,
+    'config_entries': <ANY>,
+    'configuration_url': 'http://127.0.0.1',
+    'connections': set({
+      tuple(
+        'mac',
+        'aa:bb:cc:dd:ee:ff',
+      ),
+    }),
+    'disabled_by': None,
+    'entry_type': None,
+    'hw_version': 'esp8266',
+    'id': <ANY>,
+    'identifiers': set({
+      tuple(
+        'wled',
+        'aabbccddeeff',
+      ),
+    }),
+    'is_new': False,
+    'manufacturer': 'WLED',
+    'model': 'DIY light',
+    'name': 'WLED RGB Light',
+    'name_by_user': None,
+    'suggested_area': None,
+    'sw_version': '0.8.5',
+    'via_device_id': None,
+  })
+# ---
diff --git a/tests/components/wled/test_button.py b/tests/components/wled/test_button.py
index daa6839557b..c1f3165e5bc 100644
--- a/tests/components/wled/test_button.py
+++ b/tests/components/wled/test_button.py
@@ -1,41 +1,41 @@
 """Tests for the WLED button platform."""
 from unittest.mock import MagicMock
 
-from freezegun import freeze_time
 import pytest
+from syrupy.assertion import SnapshotAssertion
 from wled import WLEDConnectionError, WLEDError
 
-from homeassistant.components.button import (
-    DOMAIN as BUTTON_DOMAIN,
-    SERVICE_PRESS,
-    ButtonDeviceClass,
-)
-from homeassistant.const import (
-    ATTR_DEVICE_CLASS,
-    ATTR_ENTITY_ID,
-    STATE_UNAVAILABLE,
-    STATE_UNKNOWN,
-    EntityCategory,
-)
+from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
+from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-pytestmark = pytest.mark.usefixtures("init_integration")
+pytestmark = [
+    pytest.mark.usefixtures("init_integration"),
+    pytest.mark.freeze_time("2021-11-04 17:37:00+01:00"),
+]
 
 
 async def test_button_restart(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    mock_wled: MagicMock,
+    snapshot: SnapshotAssertion,
 ) -> None:
     """Test the creation and values of the WLED button."""
     assert (state := hass.states.get("button.wled_rgb_light_restart"))
+    assert state == snapshot
+
+    assert (entity_entry := entity_registry.async_get(state.entity_id))
+    assert entity_entry == snapshot
+
+    assert entity_entry.device_id
+    assert (device_entry := device_registry.async_get(entity_entry.device_id))
+    assert device_entry == snapshot
+
     assert state.state == STATE_UNKNOWN
-    assert state.attributes[ATTR_DEVICE_CLASS] == ButtonDeviceClass.RESTART
-
-    assert (entry := entity_registry.async_get("button.wled_rgb_light_restart"))
-    assert entry.unique_id == "aabbccddeeff_restart"
-    assert entry.entity_category is EntityCategory.CONFIG
-
     await hass.services.async_call(
         BUTTON_DOMAIN,
         SERVICE_PRESS,
@@ -45,15 +45,11 @@ async def test_button_restart(
     assert mock_wled.reset.call_count == 1
     mock_wled.reset.assert_called_with()
 
+    assert (state := hass.states.get("button.wled_rgb_light_restart"))
+    assert state.state == "2021-11-04T16:37:00+00:00"
 
-@freeze_time("2021-11-04 17:37:00", tz_offset=-1)
-async def test_button_error(
-    hass: HomeAssistant,
-    mock_wled: MagicMock,
-) -> None:
-    """Test error handling of the WLED buttons."""
+    # Test with WLED error
     mock_wled.reset.side_effect = WLEDError
-
     with pytest.raises(HomeAssistantError, match="Invalid response from WLED API"):
         await hass.services.async_call(
             BUTTON_DOMAIN,
@@ -63,17 +59,12 @@ async def test_button_error(
         )
         await hass.async_block_till_done()
 
+    # Ensure this didn't made the entity unavailable
     assert (state := hass.states.get("button.wled_rgb_light_restart"))
-    assert state.state == "2021-11-04T16:37:00+00:00"
+    assert state.state != STATE_UNAVAILABLE
 
-
-async def test_button_connection_error(
-    hass: HomeAssistant,
-    mock_wled: MagicMock,
-) -> None:
-    """Test error handling of the WLED buttons."""
+    # Test with WLED connection error
     mock_wled.reset.side_effect = WLEDConnectionError
-
     with pytest.raises(HomeAssistantError, match="Error communicating with WLED API"):
         await hass.services.async_call(
             BUTTON_DOMAIN,
@@ -82,5 +73,6 @@ async def test_button_connection_error(
             blocking=True,
         )
 
+    # Ensure this made the entity unavailable
     assert (state := hass.states.get("button.wled_rgb_light_restart"))
     assert state.state == STATE_UNAVAILABLE