From 1ffa8fcbba94be80bd554c916145319e70920db6 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Fri, 12 Jun 2020 17:54:46 -0700
Subject: [PATCH] Platforms cleanup when adding entity fails (#36742)

---
 homeassistant/helpers/entity_platform.py |  8 +++++++
 tests/helpers/test_entity_platform.py    | 29 ++++++++++++++++++++----
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py
index 28fd83d99c1..717917f816d 100644
--- a/homeassistant/helpers/entity_platform.py
+++ b/homeassistant/helpers/entity_platform.py
@@ -319,6 +319,8 @@ class EntityPlatform:
                 await entity.async_device_update(warning=False)
             except Exception:  # pylint: disable=broad-except
                 self.logger.exception("%s: Error on device update!", self.platform_name)
+                entity.hass = None
+                entity.platform = None
                 return
 
         suggested_object_id = None
@@ -391,6 +393,8 @@ class EntityPlatform:
                     or entity.name
                     or f'"{self.platform_name} {entity.unique_id}"',
                 )
+                entity.hass = None
+                entity.platform = None
                 return
 
         # We won't generate an entity ID if the platform has already set one
@@ -416,6 +420,8 @@ class EntityPlatform:
 
         # Make sure it is valid in case an entity set the value themselves
         if not valid_entity_id(entity.entity_id):
+            entity.hass = None
+            entity.platform = None
             raise HomeAssistantError(f"Invalid entity id: {entity.entity_id}")
 
         already_exists = entity.entity_id in self.entities
@@ -431,6 +437,8 @@ class EntityPlatform:
             if entity.unique_id is not None:
                 msg += f". Platform {self.platform_name} does not generate unique IDs"
             self.logger.error(msg)
+            entity.hass = None
+            entity.platform = None
             return
 
         entity_id = entity.entity_id
diff --git a/tests/helpers/test_entity_platform.py b/tests/helpers/test_entity_platform.py
index 11dded7416f..8917e69e3ae 100644
--- a/tests/helpers/test_entity_platform.py
+++ b/tests/helpers/test_entity_platform.py
@@ -7,7 +7,7 @@ import pytest
 
 from homeassistant.const import UNIT_PERCENTAGE
 from homeassistant.core import callback
-from homeassistant.exceptions import PlatformNotReady
+from homeassistant.exceptions import HomeAssistantError, PlatformNotReady
 from homeassistant.helpers import entity_platform, entity_registry
 from homeassistant.helpers.entity import async_generate_entity_id
 from homeassistant.helpers.entity_component import (
@@ -359,6 +359,11 @@ async def test_raise_error_on_update(hass):
     assert len(updates) == 1
     assert 1 in updates
 
+    assert entity1.hass is None
+    assert entity1.platform is None
+    assert entity2.hass is not None
+    assert entity2.platform is not None
+
 
 async def test_async_remove_with_platform(hass):
     """Remove an entity from a platform."""
@@ -380,10 +385,11 @@ async def test_not_adding_duplicate_entities_with_unique_id(hass):
 
     assert len(hass.states.async_entity_ids()) == 1
 
-    await component.async_add_entities(
-        [MockEntity(name="test2", unique_id="not_very_unique")]
-    )
+    ent2 = MockEntity(name="test2", unique_id="not_very_unique")
+    await component.async_add_entities([ent2])
 
+    assert ent2.hass is None
+    assert ent2.platform is None
     assert len(hass.states.async_entity_ids()) == 1
 
 
@@ -792,6 +798,11 @@ async def test_entity_disabled_by_integration(hass):
 
     await component.async_add_entities([entity_default, entity_disabled])
 
+    assert entity_default.hass is not None
+    assert entity_default.platform is not None
+    assert entity_disabled.hass is None
+    assert entity_disabled.platform is None
+
     registry = await hass.helpers.entity_registry.async_get_registry()
 
     entry_default = registry.async_get_or_create(DOMAIN, DOMAIN, "default")
@@ -889,3 +900,13 @@ async def test_platforms_sharing_services(hass):
     assert len(entities) == 2
     assert entity1 in entities
     assert entity2 in entities
+
+
+async def test_invalid_entity_id(hass):
+    """Test specifying an invalid entity id."""
+    platform = MockEntityPlatform(hass)
+    entity = MockEntity(entity_id="invalid_entity_id")
+    with pytest.raises(HomeAssistantError):
+        await platform.async_add_entities([entity])
+    assert entity.hass is None
+    assert entity.platform is None