Use index to query available entities

detective/core.py

@@ -80,26 +80,28 @@ def perform_query(self, query):
             raise
 
     def fetch_entities(self):
-        """Fetch unique entities which have data (COUNT>0)."""
+        """Fetch entities for which we have data."""
         query = text(
             """
-            SELECT entity_id, COUNT(*)
+            SELECT entity_id
             FROM states
             GROUP BY entity_id
-            ORDER by 2 DESC
             """
         )
         response = self.perform_query(query)
-        entities = [e[0] for e in list(response)]
-        print("There are {} entities with data".format(len(entities)))
 
         # Parse the domains from the entities.
-        self._domains = list(set([e.split(".")[0] for e in entities]))
-        self._entities = {}
+        entities = {}
+        domains = set()
+
+        for [entity] in response:
+            domain = entity.split('.')[0]
+            domains.add(domain)
+            entities.setdefault(domain, []).append(entity)
 
-        # Parse entities into a dict indexed by domain.
-        for d in self._domains:
-            self._entities[d] = [e for e in entities if e.split(".")[0] == d]
+        self._domains = list(domains)
+        self._entities = entities
+        print("There are {} entities with data".format(len(entities)))
 
     def fetch_data_by_list(self, entities: List[str], limit=50000):
         """

tests/conftest.py

@@ -0,0 +1,11 @@
+from unittest.mock import patch
+
+import pytest
+
+from detective.core import HassDatabase
+
+
+@pytest.fixture
+def mock_db():
+    with patch('detective.core.create_engine'):
+        return HassDatabase('mock://db')

tests/test_core.py

@@ -1,3 +1,5 @@
+from unittest.mock import patch
+
 from detective.core import get_db_type, stripped_db_url
 
 
@@ -12,3 +14,17 @@ def test_stripped_db_url():
         'mysql://paulus@localhost'
     assert stripped_db_url('mysql://paulus:password@localhost') == \
         'mysql://paulus:***@localhost'
+
+
+def test_fetch_entities(mock_db):
+    with patch.object(mock_db, 'perform_query', return_value=[
+        ['light.kitchen'],
+        ['light.living_room'],
+        ['switch.ac'],
+    ]):
+        mock_db.fetch_entities()
+
+    assert mock_db.entities == {
+        'light': ['light.kitchen', 'light.living_room'],
+        'switch': ['switch.ac'],
+    }