Add locks to entity component

homeassistant/helpers/entity_component.py

@@ -4,6 +4,8 @@ homeassistant.helpers.entity_component
 
 Provides helpers for components that manage entities.
 """
+from threading import Lock
+
 from homeassistant.bootstrap import prepare_setup_platform
 from homeassistant.helpers import (
     generate_entity_id, config_per_platform, extract_entity_ids)
@@ -37,6 +39,7 @@ class EntityComponent(object):
         self.is_polling = False
 
         self.config = None
+        self.lock = Lock()
 
     def setup(self, config):
         """
@@ -61,8 +64,11 @@ class EntityComponent(object):
         Takes in a list of new entities. For each entity will see if it already
         exists. If not, will add it, set it up and push the first state.
         """
+        with self.lock:
             for entity in new_entities:
-            if entity is not None and entity not in self.entities.values():
+                if entity is None or entity in self.entities.values():
+                    continue
+
                 entity.hass = self.hass
 
                 if getattr(entity, 'entity_id', None) is None:
@@ -81,16 +87,26 @@ class EntityComponent(object):
             if self.group is not None:
                 self.group.update_tracked_entity_ids(self.entities.keys())
 
-        self._start_polling()
+            if self.is_polling or \
+               not any(entity.should_poll for entity
+                       in self.entities.values()):
+                return
+
+            self.is_polling = True
+
+            track_utc_time_change(
+                self.hass, self._update_entity_states,
+                second=range(0, 60, self.scan_interval))
 
     def extract_from_service(self, service):
         """
         Takes a service and extracts all known entities.
         Will return all if no entity IDs given in service.
         """
+        with self.lock:
             if ATTR_ENTITY_ID not in service.data:
-            return self.entities.values()
-        else:
+                return list(self.entities.values())
+
             return [self.entities[entity_id] for entity_id
                     in extract_entity_ids(self.hass, service)
                     if entity_id in self.entities]
@@ -99,6 +115,7 @@ class EntityComponent(object):
         """ Update the states of all the entities. """
         self.logger.info("Updating %s entities", self.domain)
 
+        with self.lock:
             for entity in self.entities.values():
                 if entity.should_poll:
                     entity.update_ha_state(True)
@@ -110,18 +127,6 @@ class EntityComponent(object):
 
         self._setup_platform(self.discovery_platforms[service], {}, info)
 
-    def _start_polling(self):
-        """ Start polling entities if necessary. """
-        if self.is_polling or \
-           not any(entity.should_poll for entity in self.entities.values()):
-            return
-
-        self.is_polling = True
-
-        track_utc_time_change(
-            self.hass, self._update_entity_states,
-            second=range(0, 60, self.scan_interval))
-
     def _setup_platform(self, platform_type, platform_config,
                         discovery_info=None):
         """ Tries to setup a platform for this component. """