diff --git a/lib/entity.class.dart b/lib/entity.class.dart
index ea13157..d4a10c2 100644
--- a/lib/entity.class.dart
+++ b/lib/entity.class.dart
@@ -1,9 +1,34 @@
 part of 'main.dart';
 
 class Entity {
+  Map _attributes;
+  String _domain;
+  String _entityId;
+  String _state;
+  String _entityPicture;
 
-  Entity() {
-    //
+  String get displayName => _attributes["friendly_name"] ?? (_attributes["name"] ?? "_");
+  String get domain => _domain;
+  String get entityId => _entityId;
+  String get state => _state;
+  String get deviceClass => _attributes["device_class"] ?? null;
+  bool get isView => (_domain == "group") && (_attributes != null ? _attributes["view"] ?? false : false);
+  bool get isGroup => _domain == "group";
+  String get icon => _attributes["icon"] ?? "";
+  bool get isOn => state == "on";
+  String get entityPicture => _attributes["entity_picture"];
+  String get unitOfMeasurement => _attributes["unit_of_measurement"] ?? "";
+  List get childEntities => _attributes["entity_id"] ?? [];
+
+  Entity(Map rawData) {
+    update(rawData);
+  }
+
+  void update(Map rawData) {
+    _attributes = rawData["attributes"] ?? {};
+    _domain = rawData["entity_id"].split(".")[0];
+    _entityId = rawData["entity_id"];
+    _state = rawData["state"];
   }
 
 }
\ No newline at end of file
diff --git a/lib/entity_collection.class.dart b/lib/entity_collection.class.dart
index ee904c3..f776848 100644
--- a/lib/entity_collection.class.dart
+++ b/lib/entity_collection.class.dart
@@ -3,33 +3,121 @@ part of 'main.dart';
 class EntityCollection {
 
   Map<String, Entity> _entities;
+  Map<String, dynamic> _uiStructure = {};
+
+  List _topBadgeDomains = ["alarm_control_panel", "binary_sensor", "device_tracker", "updater", "sun", "timer", "sensor"];
 
   EntityCollection() {
     _entities = {};
   }
 
-  void fillFromRawData(Map rawData) {
-    _entities.clear();
-    if (response["success"] == false) {
-      _statesCompleter.completeError({"errorCode": 3, "errorMessage": response["error"]["message"]});
-      return;
-    }
-    List data = response["result"];
-    TheLogger.log("Debug","Parsing ${data.length} Home Assistant entities");
-    List<String> viewsList = [];
-    data.forEach((entity) {
-      try {
-        var composedEntity = _parseEntity(entity);
+  Map<String, dynamic> get ui => _uiStructure;
 
-        if (composedEntity["attributes"] != null) {
-          if ((composedEntity["domain"] == "group") &&
-              (composedEntity["attributes"]["view"] == true)) {
-            viewsList.add(composedEntity["entity_id"]);
-          }
+  void parse(List rawData) {
+    _entities.clear();
+    _uiStructure.clear();
+
+    TheLogger.log("Debug","Parsing ${rawData.length} Home Assistant entities");
+    rawData.forEach((rawEntityData) {
+      Entity newEntity = addFromRaw(rawEntityData);
+
+      if (newEntity.isView) {
+        _uiStructure.addAll({newEntity.entityId: {}});
+      }
+    });
+
+    _createViews();
+  }
+
+  void updateState(Map rawStateData) {
+    if (isExist(rawStateData["entity_id"])) {
+      updateFromRaw(rawStateData["new_state"]);
+    } else {
+      addFromRaw(rawStateData["new_state"]);
+    }
+  }
+
+  void add(Entity entity) {
+    _entities[entity.entityId] = entity;
+  }
+
+  Entity addFromRaw(Map rawEntityData) {
+    Entity entity = Entity(rawEntityData);
+    _entities[entity.entityId] = entity;
+    return entity;
+  }
+
+  void updateFromRaw(Map rawEntityData) {
+    //TODO pass entity in this function and call update from it
+    _entities[rawEntityData["entity_id"]].update(rawEntityData);
+  }
+
+  Entity get(String entityId) {
+    return _entities[entityId];
+  }
+
+  bool isExist(String entityId) {
+    return _entities[entityId] != null;
+  }
+
+  void _createViews() async {
+    TheLogger.log("Debug","Gethering views");
+    int viewCounter = 0;
+    _uiStructure.forEach((viewId, viewStructure) { //Each view
+      try {
+        viewCounter += 1;
+        Entity viewGroupData = get(viewId);
+        if (viewGroupData != null) {
+          viewStructure["groups"] = {};
+          viewStructure["state"] = "on";
+          viewStructure["entity_id"] = viewGroupData.entityId;
+          viewStructure["badges"] = {"children": []};
+          viewStructure["attributes"] = {
+            "icon": viewGroupData.icon
+          };
+
+          viewGroupData.childEntities.forEach((
+              entityId) { //Each entity or group in view
+            Map newGroup = {};
+            if (isExist(entityId)) {
+              Entity cardOrEntityData = get(entityId);
+              if (!cardOrEntityData.isGroup) {
+                if (_topBadgeDomains.contains(cardOrEntityData.domain)) {
+                  viewStructure["badges"]["children"].add(entityId);
+                } else {
+                  String autoGroupID = "${cardOrEntityData.domain}.${cardOrEntityData.domain}$viewCounter";
+                  if (viewStructure["groups"]["$autoGroupID"] == null) {
+                    newGroup["entity_id"] = autoGroupID;
+                    newGroup["friendly_name"] = cardOrEntityData.domain;
+                    newGroup["children"] = [];
+                    newGroup["children"].add(entityId);
+                    viewStructure["groups"]["$autoGroupID"] =
+                        Map.from(newGroup);
+                  } else {
+                    viewStructure["groups"]["$autoGroupID"]["children"].add(
+                        entityId);
+                  }
+                }
+              } else {
+                newGroup["entity_id"] = entityId;
+                newGroup["friendly_name"] = cardOrEntityData.displayName;
+                newGroup["children"] = List<String>();
+                cardOrEntityData.childEntities.forEach((
+                    groupedEntityId) {
+                  newGroup["children"].add(groupedEntityId);
+                });
+                viewStructure["groups"]["$entityId"] = Map.from(newGroup);
+              }
+            } else {
+              TheLogger.log("Warning", "Unknown entity inside view: $entityId");
+            }
+          });
+        } else {
+          TheLogger.log("Warning", "No state or attributes found for view: $viewId");
         }
-        _entitiesData[entity["entity_id"]] = composedEntity;
+
       } catch (error) {
-        TheLogger.log("Error","Error parsing entity: ${entity['entity_id']}");
+        TheLogger.log("Error","Error parsing view: $viewId");
       }
     });
   }
diff --git a/lib/data_provider.class.dart b/lib/home_assistant.class.dart
similarity index 58%
rename from lib/data_provider.class.dart
rename to lib/home_assistant.class.dart
index 3549e8f..eb9f874 100644
--- a/lib/data_provider.class.dart
+++ b/lib/home_assistant.class.dart
@@ -1,35 +1,36 @@
 part of 'main.dart';
 
-class HADataProvider {
+class HomeAssistant {
   String _hassioAPIEndpoint;
   String _hassioPassword;
   String _hassioAuthType;
+
   IOWebSocketChannel _hassioChannel;
+
   int _currentMessageId = 0;
   int _statesMessageId = 0;
   int _servicesMessageId = 0;
   int _subscriptionMessageId = 0;
   int _configMessageId = 0;
-  Map _entitiesData = {};
-  Map _servicesData = {};
-  Map _uiStructure = {};
+  EntityCollection _entities;
   Map _instanceConfig = {};
+
   Completer _fetchCompleter;
   Completer _statesCompleter;
   Completer _servicesCompleter;
   Completer _configCompleter;
   Timer _fetchingTimer;
-  List _topBadgeDomains = ["alarm_control_panel", "binary_sensor", "device_tracker", "updater", "sun", "timer", "sensor"];
 
-  Map get entities => _entitiesData;
-  Map get services => _servicesData;
-  Map get uiStructure => _uiStructure;
-  Map get instanceConfig => _instanceConfig;
+  String get locationName => _instanceConfig["location_name"] ?? "";
 
-  HADataProvider(String url, String password, String authType) {
+
+  EntityCollection get entities => _entities;
+
+  HomeAssistant(String url, String password, String authType) {
     _hassioAPIEndpoint = url;
     _hassioPassword = password;
     _hassioAuthType = authType;
+    _entities = EntityCollection();
   }
 
   Future fetch() {
@@ -179,19 +180,9 @@ class HADataProvider {
   }
 
   void _handleEntityStateChange(Map eventData) {
-    //TheLogger.log("Debug", "Parsing new state for ${eventData['entity_id']}");
-    if (eventData["new_state"] == null) {
-      TheLogger.log("Error", "No new_state found");
-    } else {
-      var parsedEntityData = _parseEntity(eventData["new_state"]);
-      String entityId = parsedEntityData["entity_id"];
-      if (_entitiesData[entityId] == null) {
-        _entitiesData[entityId] = parsedEntityData;
-      } else {
-        _entitiesData[entityId].addAll(parsedEntityData);
-      }
-      eventBus.fire(new StateChangedEvent(eventData["entity_id"]));
-    }
+    TheLogger.log("Debug", "Parsing new state for ${eventData['entity_id']}");
+    _entities.updateState(eventData);
+    eventBus.fire(new StateChangedEvent(eventData["entity_id"]));
   }
 
   void _parseConfig(Map data) {
@@ -204,7 +195,8 @@ class HADataProvider {
   }
 
   void _parseServices(response) {
-    if (response["success"] == false) {
+    _servicesCompleter.complete();
+    /*if (response["success"] == false) {
       _servicesCompleter.completeError({"errorCode": 4, "errorMessage": response["error"]["message"]});
       return;
     }
@@ -215,7 +207,7 @@ class HADataProvider {
       data.forEach((domain, services) {
         result[domain] = Map.from(services);
         services.forEach((serviceName, serviceData) {
-          if (_entitiesData["$domain.$serviceName"] != null) {
+          if (_entitiesData.isExist("$domain.$serviceName")) {
             result[domain].remove(serviceName);
           }
         });
@@ -223,116 +215,20 @@ class HADataProvider {
       _servicesData = result;
       _servicesCompleter.complete();
     } catch (e) {
-      //TODO hadle it properly
       TheLogger.log("Error","Error parsing services. But they are not used :-)");
       _servicesCompleter.complete();
-    }
+    }*/
   }
 
   void _parseEntities(response) async {
-    _entitiesData.clear();
-    _uiStructure.clear();
     if (response["success"] == false) {
       _statesCompleter.completeError({"errorCode": 3, "errorMessage": response["error"]["message"]});
       return;
     }
-    List data = response["result"];
-    TheLogger.log("Debug","Parsing ${data.length} Home Assistant entities");
-    List<String> viewsList = [];
-    data.forEach((entity) {
-      try {
-        var composedEntity = _parseEntity(entity);
-
-        if (composedEntity["attributes"] != null) {
-          if ((composedEntity["domain"] == "group") &&
-              (composedEntity["attributes"]["view"] == true)) {
-            viewsList.add(composedEntity["entity_id"]);
-          }
-        }
-        _entitiesData[entity["entity_id"]] = composedEntity;
-      } catch (error) {
-        TheLogger.log("Error","Error parsing entity: ${entity['entity_id']}");
-      }
-    });
-
-    //Gethering information for UI
-    TheLogger.log("Debug","Gethering views");
-    int viewCounter = 0;
-    viewsList.forEach((viewId) { //Each view
-      try {
-        Map viewStructure = {};
-        viewCounter += 1;
-        var viewGroupData = _entitiesData[viewId];
-        if ((viewGroupData != null) && (viewGroupData["attributes"] != null)) {
-          viewStructure["groups"] = {};
-          viewStructure["state"] = "on";
-          viewStructure["entity_id"] = viewGroupData["entity_id"];
-          viewStructure["badges"] = {"children": []};
-          viewStructure["attributes"] = {
-              "icon": viewGroupData["attributes"]["icon"]
-            };
-
-          viewGroupData["attributes"]["entity_id"].forEach((
-              entityId) { //Each entity or group in view
-            Map newGroup = {};
-            if (_entitiesData[entityId] != null) {
-              Map cardOrEntityData = _entitiesData[entityId];
-              String domain = cardOrEntityData["domain"];
-              if (domain != "group") {
-                if (_topBadgeDomains.contains(domain)) {
-                  viewStructure["badges"]["children"].add(entityId);
-                } else {
-                  String autoGroupID = "$domain.$domain$viewCounter";
-                  if (viewStructure["groups"]["$autoGroupID"] == null) {
-                    newGroup["entity_id"] = "$domain.$domain$viewCounter";
-                    newGroup["friendly_name"] = "$domain";
-                    newGroup["children"] = [];
-                    newGroup["children"].add(entityId);
-                    viewStructure["groups"]["$autoGroupID"] =
-                        Map.from(newGroup);
-                  } else {
-                    viewStructure["groups"]["$autoGroupID"]["children"].add(
-                        entityId);
-                  }
-                }
-              } else {
-                if (cardOrEntityData["attributes"] != null) {
-                  newGroup["entity_id"] = entityId;
-                  newGroup["friendly_name"] = cardOrEntityData['attributes']['friendly_name'] ?? "";
-                  newGroup["children"] = List<String>();
-                  cardOrEntityData["attributes"]["entity_id"].forEach((
-                      groupedEntityId) {
-                    newGroup["children"].add(groupedEntityId);
-                  });
-                  viewStructure["groups"]["$entityId"] = Map.from(newGroup);
-                } else {
-                  TheLogger.log("Warning", "Group has no attributes to build a card: $entityId");
-                }
-              }
-            } else {
-              TheLogger.log("Warning", "Unknown entity inside view: $entityId");
-            }
-          });
-          _uiStructure[viewId.split(".")[1]] = viewStructure;
-        } else {
-          TheLogger.log("Warning", "No state or attributes found for view: $viewId");
-        }
-
-      } catch (error) {
-        TheLogger.log("Error","Error parsing view: $viewId");
-      }
-    });
+    _entities.parse(response["result"]);
     _statesCompleter.complete();
   }
 
-  Map _parseEntity(rawData) {
-    var composedEntity = Map.from(rawData);
-    String entityDomain = rawData["entity_id"].split(".")[0];
-    composedEntity["display_name"] = "${rawData["attributes"]!=null ? rawData["attributes"]["friendly_name"] ?? rawData["attributes"]["name"] : "_"}";
-    composedEntity["domain"] = entityDomain;
-    return composedEntity;
-  }
-
   Future callService(String domain, String service, String entity_id) {
     var sendCompleter = Completer();
     //TODO: Send service call timeout timer. Should be removed after #21 fix
diff --git a/lib/main.dart b/lib/main.dart
index f158c5f..d15d821 100644
--- a/lib/main.dart
+++ b/lib/main.dart
@@ -12,7 +12,7 @@ import 'package:url_launcher/url_launcher.dart';
 import 'package:flutter/services.dart';
 
 part 'settings.page.dart';
-part 'data_provider.class.dart';
+part 'home_assistant.class.dart';
 part 'log.page.dart';
 part 'utils.class.dart';
 part 'mdi.class.dart';
@@ -37,6 +37,9 @@ void main() {
     runApp(new HAClientApp());
   }, onError: (error, stack) {
     TheLogger.log("Global error", "$error");
+    if (TheLogger.isInDebugMode) {
+      debugPrint("$stack");
+    }
   });
 }
 
@@ -69,10 +72,10 @@ class MainPage extends StatefulWidget {
 }
 
 class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
-  HADataProvider _dataModel;
-  Map _entitiesData;
+  HomeAssistant _homeAssistant;
+  EntityCollection _entities;
   Map _uiStructure;
-  Map _instanceConfig;
+  //Map _instanceConfig;
   int _uiViewsCount = 0;
   String _instanceHost;
   int _errorCodeToBeShown = 0;
@@ -129,18 +132,18 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
         _errorCodeToBeShown = 5;
       });
     } else {
-      if (_dataModel != null) _dataModel.closeConnection();
+      if (_homeAssistant != null) _homeAssistant.closeConnection();
       _createConnection(apiEndpoint, apiPassword, authType);
     }
   }
 
   _createConnection(String apiEndpoint, String apiPassword, String authType) {
-    _dataModel = HADataProvider(apiEndpoint, apiPassword, authType);
+    _homeAssistant = HomeAssistant(apiEndpoint, apiPassword, authType);
     _refreshData();
     if (_stateSubscription != null) _stateSubscription.cancel();
     _stateSubscription = eventBus.on<StateChangedEvent>().listen((event) {
       setState(() {
-        _entitiesData = _dataModel.entities;
+        _entities = _homeAssistant.entities;
       });
     });
   }
@@ -150,12 +153,12 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
       _isLoading = true;
     });
     _errorCodeToBeShown = 0;
-    if (_dataModel != null) {
-      await _dataModel.fetch().then((result) {
+    if (_homeAssistant != null) {
+      await _homeAssistant.fetch().then((result) {
         setState(() {
-          _instanceConfig = _dataModel.instanceConfig;
-          _entitiesData = _dataModel.entities;
-          _uiStructure = _dataModel.uiStructure;
+          //_instanceConfig = _homeAssistant.instanceConfig;
+          _entities = _homeAssistant.entities;
+          _uiStructure = _entities.ui;
           _uiViewsCount = _uiStructure.length;
           _isLoading = false;
         });
@@ -177,7 +180,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
     setState(() {
       _isLoading = true;
     });
-    _dataModel.callService(domain, service, entityId).then((r) {
+    _homeAssistant.callService(domain, service, entityId).then((r) {
       setState(() {
         _isLoading = false;
       });
@@ -186,7 +189,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
 
   List<Widget> _buildViews() {
     List<Widget> result = [];
-    if ((_entitiesData != null) && (_uiStructure != null)) {
+    if ((_entities != null) && (_uiStructure != null)) {
       _uiStructure.forEach((viewId, structure) {
         result.add(
             RefreshIndicator(
@@ -229,7 +232,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
   List<Widget> _buildBadges(List ids) {
     List<Widget> result = [];
     ids.forEach((entityId) {
-      var data = _entitiesData[entityId];
+      var data = _entities.get(entityId);
       if (data != null) {
         result.add(
           _buildSingleBadge(data)
@@ -239,14 +242,14 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
     return result;
   }
 
-  Widget _buildSingleBadge(data) {
+  Widget _buildSingleBadge(Entity data) {
     double iconSize = 26.0;
     Widget badgeIcon;
     String badgeTextValue;
-    Color iconColor = _badgeColors[data["domain"]] ?? _badgeColors["default"];
-    switch (data["domain"]) {
+    Color iconColor = _badgeColors[data.domain] ?? _badgeColors["default"];
+    switch (data.domain) {
       case "sun": {
-        badgeIcon = data["state"] == "below_horizon" ?
+        badgeIcon = data.state == "below_horizon" ?
           Icon(
             MaterialDesignIcons.createIconDataFromIconCode(0xf0dc),
             size: iconSize,
@@ -258,10 +261,10 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
         break;
       }
       case "sensor": {
-        badgeTextValue = data["attributes"]["unit_of_measurement"];
+        badgeTextValue = data.unitOfMeasurement;
         badgeIcon = Center(
           child: Text(
-            "${data['state']}",
+            "${data.state}",
             overflow: TextOverflow.fade,
             softWrap: false,
             textAlign: TextAlign.center,
@@ -272,7 +275,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
       }
       case "device_tracker": {
         badgeIcon = MaterialDesignIcons.createIconFromEntityData(data, iconSize,Colors.black);
-        badgeTextValue = data["state"];
+        badgeTextValue = data.state;
         break;
       }
       default: {
@@ -337,7 +340,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
         Container(
           width: 60.0,
           child: Text(
-            "${data['display_name']}",
+            "${data.displayName}",
             textAlign: TextAlign.center,
             softWrap: true,
             maxLines: 2,
@@ -378,14 +381,14 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
   List<Widget> _buildCardBody(List ids) {
     List<Widget> entities = [];
     ids.forEach((id) {
-      var data = _entitiesData[id];
+      var data = _entities.get(id);
       if (data != null) {
         entities.add(new ListTile(
-          leading: MaterialDesignIcons.createIconFromEntityData(data, 28.0, _stateIconColors[data["state"]] ?? Colors.blueGrey),
+          leading: MaterialDesignIcons.createIconFromEntityData(data, 28.0, _stateIconColors[data.state] ?? Colors.blueGrey),
           //subtitle: Text("${data['entity_id']}"),
           trailing: _buildEntityActionWidget(data),
           title: Text(
-            "${data["display_name"]}",
+            "${data.displayName}",
             overflow: TextOverflow.fade,
             softWrap: false,
           ),
@@ -395,21 +398,23 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
     return entities;
   }
 
-  Widget _buildEntityActionWidget(data) {
-    String entityId = data["entity_id"];
+  Widget _buildEntityActionWidget(Entity entity) {
+    String entityId = entity.entityId;
     Widget result;
-    switch (data["domain"]) {
+    switch (entity.domain) {
       case "automation":
       case "switch":
       case "light": {
         result = Switch(
-          value: (data["state"] == "on"),
+          value: entity.isOn,
           onChanged: ((state) {
             _callService(
-                data["domain"], state ? "turn_on" : "turn_off", entityId);
-            setState(() {
-              _entitiesData[entityId]["state"] = state ? "on" : "off";
-            });
+                entity.domain, state ? "turn_on" : "turn_off", entityId
+            );
+            //TODO remove after checking if state will chenge without setState but after socket event
+            /*setState(() {
+              _entities[entityId]["state"] = state ? "on" : "off";
+            });*/
           }),
         );
         break;
@@ -421,7 +426,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
           width: 60.0,
           child: FlatButton(
             onPressed: (() {
-              _callService(data["domain"], "turn_on", entityId);
+              _callService(entity.domain, "turn_on", entityId);
             }),
             child: Text(
               "Run",
@@ -437,7 +442,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
         result = Padding(
           padding: EdgeInsets.fromLTRB(0.0, 0.0, 16.0, 0.0),
           child: Text(
-            "${data["state"]}${(data["attributes"] != null && data["attributes"]["unit_of_measurement"] != null) ? data["attributes"]["unit_of_measurement"] : ''}",
+            "${entity.state}${entity.unitOfMeasurement}",
             textAlign: TextAlign.right,
             style: new TextStyle(
               fontSize: 16.0,
@@ -457,11 +462,11 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
 
   List<Tab> buildUIViewTabs() {
     List<Tab> result = [];
-    if ((_entitiesData != null) && (_uiStructure != null)) {
+    if ((_entities != null) && (_uiStructure != null)) {
       _uiStructure.forEach((viewId, structure) {
         result.add(
             Tab(
-                icon: MaterialDesignIcons.createIconFromEntityData(structure, 24.0, null)
+                icon: MaterialDesignIcons.createIconFromEntityData(_entities.get(viewId), 24.0, null)
             )
         );
       });
@@ -471,7 +476,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
 
   Widget _buildAppTitle() {
     Row titleRow = Row(
-      children: [Text(_instanceConfig != null ? _instanceConfig["location_name"] : "")],
+      children: [Text(_homeAssistant != null ? _homeAssistant.locationName : "")],
     );
     if (_isLoading) {
       titleRow.children.add(Padding(
@@ -490,7 +495,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
       child: ListView(
         children: <Widget>[
           new UserAccountsDrawerHeader(
-            accountName: Text(_instanceConfig != null ? _instanceConfig["location_name"] : "Unknown"),
+            accountName: Text(_homeAssistant != null ? _homeAssistant.locationName : "Unknown"),
             accountEmail: Text(_instanceHost ?? "Not configured"),
             currentAccountPicture: new Image.asset('images/hassio-192x192.png'),
           ),
@@ -638,7 +643,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
   Widget build(BuildContext context) {
     _checkShowInfo(context);
     // This method is rerun every time setState is called.
-    if (_entitiesData == null) {
+    if (_entities == null) {
       return _buildScaffold(true);
     } else {
       return DefaultTabController(
@@ -653,7 +658,7 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
     WidgetsBinding.instance.removeObserver(this);
     if (_stateSubscription != null) _stateSubscription.cancel();
     if (_settingsSubscription != null) _settingsSubscription.cancel();
-    _dataModel.closeConnection();
+    _homeAssistant.closeConnection();
     super.dispose();
   }
 }
diff --git a/lib/mdi.class.dart b/lib/mdi.class.dart
index 9df398d..5aecc00 100644
--- a/lib/mdi.class.dart
+++ b/lib/mdi.class.dart
@@ -2868,30 +2868,26 @@ class MaterialDesignIcons {
     "mdi:blank": 0xf68c
   };
 
-  static Widget createIconFromEntityData(Map data, double size, Color color) {
-    if ((data["attributes"] != null) && (data["attributes"]["entity_picture"] != null)) {
+  static Widget createIconFromEntityData(Entity data, double size, Color color) {
+    if (data.entityPicture != null) {
       if (homeAssistantWebHost != null) {
         return CircleAvatar(
           backgroundColor: Colors.white,
           backgroundImage: CachedNetworkImageProvider(
-            "$homeAssistantWebHost${data["attributes"]["entity_picture"]}",
+            "$homeAssistantWebHost${data.entityPicture}",
           ),
         );
       } else {
         return Container(width: 0.0, height: 0.0);
       }
     } else {
-      String iconName = data["attributes"] != null
-          ? data["attributes"]["icon"]
-          : null;
+      String iconName = data.icon;
       int iconCode = 0;
-      if (iconName != null) {
+      if (iconName.length > 0) {
         iconCode = getIconCodeByIconName(iconName);
       } else {
-        iconCode = getDefaultIconByEntityId(data["entity_id"],
-            data["attributes"] != null
-                ? data["attributes"]["device_class"]
-                : null, data["state"]); //
+        iconCode = getDefaultIconByEntityId(data.entityId,
+            data.deviceClass, data.state); //
       }
       return Icon(
         IconData(iconCode, fontFamily: 'Material Design Icons'),