diff --git a/package.json b/package.json
index ec9f9c6..fcb59d3 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "battery-state-card",
-  "version": "3.2.0",
+  "version": "3.2.1",
   "description": "Battery State card for Home Assistant",
   "main": "dist/battery-state-card.js",
   "author": "Max Chodorowski",
diff --git a/src/entity-fields/battery-level.ts b/src/entity-fields/battery-level.ts
index a1f3db5..4d475c7 100644
--- a/src/entity-fields/battery-level.ts
+++ b/src/entity-fields/battery-level.ts
@@ -50,8 +50,8 @@ export const getBatteryLevel = (config: IBatteryEntityConfig, hass: HomeAssistan
     }
     else {
         const candidates: (string | number | undefined)[] = [
-            config.non_battery_entity ? null: entityData.attributes.battery_level,
-            config.non_battery_entity ? null: entityData.attributes.battery,
+            config.non_battery_entity ? null: entityData.attributes?.battery_level,
+            config.non_battery_entity ? null: entityData.attributes?.battery,
             entityData.state
         ];
 
diff --git a/src/entity-fields/get-name.ts b/src/entity-fields/get-name.ts
index 92c842b..a55c59c 100644
--- a/src/entity-fields/get-name.ts
+++ b/src/entity-fields/get-name.ts
@@ -15,11 +15,7 @@ export const getName = (config: IBatteryEntityConfig, hass: HomeAssistant | unde
         return proc.process(config.name);
     }
 
-    if (!hass) {
-        return config.entity;
-    }
-
-    let name = hass.states[config.entity]?.attributes.friendly_name;
+    let name = entityData?.attributes?.friendly_name;
 
     // when we have failed to get the name we just return entity id
     if (!name) {
diff --git a/test/helpers.ts b/test/helpers.ts
index 4ca4b64..86141f3 100644
--- a/test/helpers.ts
+++ b/test/helpers.ts
@@ -195,6 +195,12 @@ export class HomeAssistantMock<T extends LovelaceCard<any>> {
                 return entity;
             },
             setAttributes: (attribs: IEntityAttributes) => {
+
+                if (attribs === null) {
+                    this.hass.states[entity.entity_id].attributes = <any>undefined;
+                    return entity;
+                }
+
                 this.hass.states[entity.entity_id].attributes = {
                     ...this.hass.states[entity.entity_id].attributes,
                     ...attribs
@@ -242,7 +248,7 @@ interface IEntityMock {
     readonly entity_id: string;
     readonly state: string;
     setState(state: string): IEntityMock;
-    setAttributes(attribs: IEntityAttributes): IEntityMock;
+    setAttributes(attribs: IEntityAttributes | null): IEntityMock;
     setLastUpdated(val: string): void;
     setLastChanged(val: string): void;
     setProperty<K extends keyof HaEntityPropertyToTypeMap>(name: K, val: HaEntityPropertyToTypeMap[K]): void;
diff --git a/test/other/entity-fields/battery-level.test.ts b/test/other/entity-fields/battery-level.test.ts
index b1e45fc..9bfd484 100644
--- a/test/other/entity-fields/battery-level.test.ts
+++ b/test/other/entity-fields/battery-level.test.ts
@@ -12,6 +12,18 @@ describe("Battery level", () => {
         expect(unit).toBe("%");
     });
 
+    test("doen't throw exception when attributes are not set on entity", () => {
+        const hassMock = new HomeAssistantMock(true);
+        const entity = hassMock.addEntity("Mocked entity", "45", { battery_state: "45" });
+        entity.setAttributes(null);
+
+        const { state, level, unit } = getBatteryLevel({ entity: "mocked_entity" }, hassMock.hass, hassMock.hass.states["mocked_entity"]);
+        
+        expect(level).toBe(45);
+        expect(state).toBe("45");
+        expect(unit).toBe("%")
+    });
+
     test("is 'Unknown' when entity not found and no localized string", () => {
         const hassMock = new HomeAssistantMock(true);
         hassMock.hass.localize = () => <string><unknown>null;
diff --git a/test/other/entity-fields/get-name.test.ts b/test/other/entity-fields/get-name.test.ts
index 1b6446e..78db2da 100644
--- a/test/other/entity-fields/get-name.test.ts
+++ b/test/other/entity-fields/get-name.test.ts
@@ -15,12 +15,22 @@ describe("Get name", () => {
         expect(name).toBe("sensor.my_entity_id");
     });
 
-    test("returns name from friendly_name attribute of the entity", () => {
+    test("doesn't throw exception when attributes property is missing", () => {
         const hassMock = new HomeAssistantMock(true);
-        hassMock.addEntity("My entity", "45", { friendly_name: "My entity name" });
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: "My entity name" });
+        entity.setAttributes(null);
 
         let name = getName({ entity: "my_entity" }, hassMock.hass, {});
 
+        expect(name).toBe("my_entity");
+    });
+
+    test("returns name from friendly_name attribute of the entity", () => {
+        const hassMock = new HomeAssistantMock(true);
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: "My entity name" });
+
+        let name = getName({ entity: "my_entity" }, hassMock.hass, hassMock.hass.states[entity.entity_id]);
+
         expect(name).toBe("My entity name");
     });
 
@@ -33,18 +43,18 @@ describe("Get name", () => {
 
     test("returns entity id when entity doesn't have a friendly_name attribute", () => {
         const hassMock = new HomeAssistantMock(true);
-        hassMock.addEntity("My entity", "45", { friendly_name: undefined });
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: undefined });
 
-        let name = getName({ entity: "my_entity" }, hassMock.hass, {});
+        let name = getName({ entity: "my_entity" }, hassMock.hass, hassMock.hass.states[entity.entity_id]);
 
         expect(name).toBe("my_entity");
     });
 
     test("doesn't throw when name is empty", () => {
         const hassMock = new HomeAssistantMock(true);
-        hassMock.addEntity("My entity", "45", { friendly_name: "Battery" });
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: "Battery" });
 
-        let name = getName({ entity: "my_entity", bulk_rename: [{ from: "Battery", to: "" }] }, hassMock.hass, {});
+        let name = getName({ entity: "my_entity", bulk_rename: [{ from: "Battery", to: "" }] }, hassMock.hass, hassMock.hass.states[entity.entity_id]);
 
         expect(name).toBe("");
     });
@@ -57,9 +67,9 @@ describe("Get name", () => {
         ]
     )("returns renamed entity name", (entityName: string, renameRules: IConvert | IConvert[], expectedResult: string) => {
         const hassMock = new HomeAssistantMock(true);
-        hassMock.addEntity("My entity", "45", { friendly_name: entityName });
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: entityName });
 
-        let name = getName({ entity: "my_entity", bulk_rename: renameRules }, hassMock.hass, {});
+        let name = getName({ entity: "my_entity", bulk_rename: renameRules }, hassMock.hass, hassMock.hass.states[entity.entity_id]);
 
         expect(name).toBe(expectedResult);
     });
@@ -72,9 +82,9 @@ describe("Get name", () => {
         ]
     )("regex", (entityName: string, renameRules: IConvert | IConvert[], expectedResult: string) => {
         const hassMock = new HomeAssistantMock(true);
-        hassMock.addEntity("My entity", "45", { friendly_name: entityName });
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: entityName });
 
-        let name = getName({ entity: "my_entity", bulk_rename: renameRules }, hassMock.hass, {});
+        let name = getName({ entity: "my_entity", bulk_rename: renameRules }, hassMock.hass, hassMock.hass.states[entity.entity_id]);
 
         expect(name).toBe(expectedResult);
     });
@@ -102,9 +112,9 @@ describe("Get name", () => {
         ["battery kitchen6", undefined, "Battery kitchen6"],
     ])("KString in the name", (entityName: string, renameRules: IBulkRename | IConvert | IConvert[] | undefined, expectedResult: string) => {
         const hassMock = new HomeAssistantMock(true);
-        hassMock.addEntity("My entity", "45", { friendly_name: entityName });
+        const entity = hassMock.addEntity("My entity", "45", { friendly_name: entityName });
 
-        let result = getName({entity: "my_entity", bulk_rename: renameRules}, hassMock.hass, {});
+        let result = getName({entity: "my_entity", bulk_rename: renameRules}, hassMock.hass, hassMock.hass.states[entity.entity_id]);
         expect(result).toBe(expectedResult);
     })
 });
\ No newline at end of file