diff --git a/src/common/entity/context/get_area_context.ts b/src/common/entity/context/get_area_context.ts
index 88927f87df..69debf4609 100644
--- a/src/common/entity/context/get_area_context.ts
+++ b/src/common/entity/context/get_area_context.ts
@@ -6,34 +6,15 @@ interface AreaContext {
   area: AreaRegistryEntry | null;
   floor: FloorRegistryEntry | null;
 }
-
-/**
- * Retrieves the context of a specific area, including its associated area registry entry
- * and floor registry entry, if available.
- *
- * @param areaId - The unique identifier of the area to retrieve context for.
- * @param hass - The Home Assistant instance containing area and floor registry data.
- * @returns An object containing the area registry entry and the associated floor registry entry,
- *          or `null` values if the area or floor is not found.
- */
 export const getAreaContext = (
-  areaId: string,
+  area: AreaRegistryEntry,
   hass: HomeAssistant
 ): AreaContext => {
-  const area = (hass.areas[areaId] as AreaRegistryEntry | undefined) || null;
-
-  if (!area) {
-    return {
-      area: null,
-      floor: null,
-    };
-  }
-
-  const floorId = area?.floor_id;
-  const floor = floorId ? hass.floors[floorId] : null;
+  const floorId = area.floor_id;
+  const floor = floorId ? hass.floors[floorId] : undefined;
 
   return {
     area: area,
-    floor: floor,
+    floor: floor || null,
   };
 };
diff --git a/src/common/entity/context/get_device_context.ts b/src/common/entity/context/get_device_context.ts
new file mode 100644
index 0000000000..ae690bc334
--- /dev/null
+++ b/src/common/entity/context/get_device_context.ts
@@ -0,0 +1,26 @@
+import type { AreaRegistryEntry } from "../../../data/area_registry";
+import type { DeviceRegistryEntry } from "../../../data/device_registry";
+import type { FloorRegistryEntry } from "../../../data/floor_registry";
+import type { HomeAssistant } from "../../../types";
+
+interface DeviceContext {
+  device: DeviceRegistryEntry;
+  area: AreaRegistryEntry | null;
+  floor: FloorRegistryEntry | null;
+}
+
+export const getDeviceContext = (
+  device: DeviceRegistryEntry,
+  hass: HomeAssistant
+): DeviceContext => {
+  const areaId = device.area_id;
+  const area = areaId ? hass.areas[areaId] : undefined;
+  const floorId = area?.floor_id;
+  const floor = floorId ? hass.floors[floorId] : undefined;
+
+  return {
+    device: device,
+    area: area || null,
+    floor: floor || null,
+  };
+};
diff --git a/src/common/entity/context/get_entity_context.ts b/src/common/entity/context/get_entity_context.ts
index ee7d35b701..0cdab2b617 100644
--- a/src/common/entity/context/get_entity_context.ts
+++ b/src/common/entity/context/get_entity_context.ts
@@ -1,6 +1,11 @@
+import type { HassEntity } from "home-assistant-js-websocket";
 import type { AreaRegistryEntry } from "../../../data/area_registry";
 import type { DeviceRegistryEntry } from "../../../data/device_registry";
-import type { EntityRegistryDisplayEntry } from "../../../data/entity_registry";
+import type {
+  EntityRegistryDisplayEntry,
+  EntityRegistryEntry,
+  ExtEntityRegistryEntry,
+} from "../../../data/entity_registry";
 import type { FloorRegistryEntry } from "../../../data/floor_registry";
 import type { HomeAssistant } from "../../../types";
 
@@ -11,27 +16,15 @@ interface EntityContext {
   floor: FloorRegistryEntry | null;
 }
 
-/**
- * Retrieves the context of an entity, including its associated device, area, and floor.
- *
- * @param entityId - The unique identifier of the entity to retrieve the context for.
- * @param hass - The Home Assistant object containing the registry data for entities, devices, areas, and floors.
- * @returns An object containing the entity, its associated device, area, and floor, or `null` for each if not found.
- *
- * The returned `EntityContext` object includes:
- * - `entity`: The entity registry entry, or `null` if the entity is not found.
- * - `device`: The device registry entry associated with the entity, or `null` if not found.
- * - `area`: The area registry entry associated with the entity or device, or `null` if not found.
- * - `floor`: The floor registry entry associated with the area, or `null` if not found.
- */
 export const getEntityContext = (
-  entityId: string,
+  stateObj: HassEntity,
   hass: HomeAssistant
 ): EntityContext => {
-  const entity =
-    (hass.entities[entityId] as EntityRegistryDisplayEntry | undefined) || null;
+  const entry = hass.entities[stateObj.entity_id] as
+    | EntityRegistryDisplayEntry
+    | undefined;
 
-  if (!entity) {
+  if (!entry) {
     return {
       entity: null,
       device: null,
@@ -39,18 +32,28 @@ export const getEntityContext = (
       floor: null,
     };
   }
+  return getEntityEntryContext(entry, hass);
+};
 
-  const deviceId = entity?.device_id;
-  const device = deviceId ? hass.devices[deviceId] : null;
-  const areaId = entity?.area_id || device?.area_id;
-  const area = areaId ? hass.areas[areaId] : null;
+export const getEntityEntryContext = (
+  entry:
+    | EntityRegistryDisplayEntry
+    | EntityRegistryEntry
+    | ExtEntityRegistryEntry,
+  hass: HomeAssistant
+): EntityContext => {
+  const entity = hass.entities[entry.entity_id];
+  const deviceId = entry?.device_id;
+  const device = deviceId ? hass.devices[deviceId] : undefined;
+  const areaId = entry?.area_id || device?.area_id;
+  const area = areaId ? hass.areas[areaId] : undefined;
   const floorId = area?.floor_id;
-  const floor = floorId ? hass.floors[floorId] : null;
+  const floor = floorId ? hass.floors[floorId] : undefined;
 
   return {
     entity: entity,
-    device: device,
-    area: area,
-    floor: floor,
+    device: device || null,
+    area: area || null,
+    floor: floor || null,
   };
 };
diff --git a/src/common/entity/entity_filter.ts b/src/common/entity/entity_filter.ts
index ae30422cc7..cef8cc1638 100644
--- a/src/common/entity/entity_filter.ts
+++ b/src/common/entity/entity_filter.ts
@@ -60,7 +60,7 @@ export const generateEntityFilter = (
       }
     }
 
-    const { area, floor, device, entity } = getEntityContext(entityId, hass);
+    const { area, floor, device, entity } = getEntityContext(stateObj, hass);
 
     if (entity && entity.hidden) {
       return false;
diff --git a/src/common/entity/get_area_context.ts b/src/common/entity/get_area_context.ts
deleted file mode 100644
index 0d14290ff3..0000000000
--- a/src/common/entity/get_area_context.ts
+++ /dev/null
@@ -1,18 +0,0 @@
-import type { AreaRegistryEntry } from "../../data/area_registry";
-import type { FloorRegistryEntry } from "../../data/floor_registry";
-import type { HomeAssistant } from "../../types";
-
-interface AreaContext {
-  floor: FloorRegistryEntry | null;
-}
-export const getAreaContext = (
-  area: AreaRegistryEntry,
-  hass: HomeAssistant
-): AreaContext => {
-  const floorId = area.floor_id;
-  const floor = floorId ? hass.floors[floorId] : null;
-
-  return {
-    floor: floor,
-  };
-};
diff --git a/src/common/entity/get_device_context.ts b/src/common/entity/get_device_context.ts
deleted file mode 100644
index c25f551350..0000000000
--- a/src/common/entity/get_device_context.ts
+++ /dev/null
@@ -1,24 +0,0 @@
-import type { AreaRegistryEntry } from "../../data/area_registry";
-import type { DeviceRegistryEntry } from "../../data/device_registry";
-import type { FloorRegistryEntry } from "../../data/floor_registry";
-import type { HomeAssistant } from "../../types";
-
-interface DeviceContext {
-  area: AreaRegistryEntry | null;
-  floor: FloorRegistryEntry | null;
-}
-
-export const getDeviceContext = (
-  device: DeviceRegistryEntry,
-  hass: HomeAssistant
-): DeviceContext => {
-  const areaId = device.area_id;
-  const area = areaId ? hass.areas[areaId] : null;
-  const floorId = area?.floor_id;
-  const floor = floorId ? hass.floors[floorId] : null;
-
-  return {
-    area: area,
-    floor: floor,
-  };
-};
diff --git a/src/common/entity/get_entity_context.ts b/src/common/entity/get_entity_context.ts
deleted file mode 100644
index a3d7b8a05b..0000000000
--- a/src/common/entity/get_entity_context.ts
+++ /dev/null
@@ -1,55 +0,0 @@
-import type { HassEntity } from "home-assistant-js-websocket";
-import type { AreaRegistryEntry } from "../../data/area_registry";
-import type { DeviceRegistryEntry } from "../../data/device_registry";
-import type {
-  EntityRegistryDisplayEntry,
-  EntityRegistryEntry,
-  ExtEntityRegistryEntry,
-} from "../../data/entity_registry";
-import type { FloorRegistryEntry } from "../../data/floor_registry";
-import type { HomeAssistant } from "../../types";
-
-interface EntityContext {
-  device: DeviceRegistryEntry | null;
-  area: AreaRegistryEntry | null;
-  floor: FloorRegistryEntry | null;
-}
-
-export const getEntityContext = (
-  stateObj: HassEntity,
-  hass: HomeAssistant
-): EntityContext => {
-  const entry = hass.entities[stateObj.entity_id] as
-    | EntityRegistryDisplayEntry
-    | undefined;
-
-  if (!entry) {
-    return {
-      device: null,
-      area: null,
-      floor: null,
-    };
-  }
-  return getEntityEntryContext(entry, hass);
-};
-
-export const getEntityEntryContext = (
-  entry:
-    | EntityRegistryDisplayEntry
-    | EntityRegistryEntry
-    | ExtEntityRegistryEntry,
-  hass: HomeAssistant
-): EntityContext => {
-  const deviceId = entry?.device_id;
-  const device = deviceId ? hass.devices[deviceId] : null;
-  const areaId = entry?.area_id || device?.area_id;
-  const area = areaId ? hass.areas[areaId] : null;
-  const floorId = area?.floor_id;
-  const floor = floorId ? hass.floors[floorId] : null;
-
-  return {
-    device: device,
-    area: area,
-    floor: floor,
-  };
-};
diff --git a/src/components/entity/ha-entity-picker.ts b/src/components/entity/ha-entity-picker.ts
index 49b5351f28..85b908e188 100644
--- a/src/components/entity/ha-entity-picker.ts
+++ b/src/components/entity/ha-entity-picker.ts
@@ -13,7 +13,7 @@ import { computeDeviceName } from "../../common/entity/compute_device_name";
 import { computeDomain } from "../../common/entity/compute_domain";
 import { computeEntityName } from "../../common/entity/compute_entity_name";
 import { computeStateName } from "../../common/entity/compute_state_name";
-import { getEntityContext } from "../../common/entity/get_entity_context";
+import { getEntityContext } from "../../common/entity/context/get_entity_context";
 import { caseInsensitiveStringCompare } from "../../common/string/compare";
 import { computeRTL } from "../../common/util/compute_rtl";
 import { domainToName } from "../../data/integration";
diff --git a/src/components/entity/ha-statistic-picker.ts b/src/components/entity/ha-statistic-picker.ts
index 165ca1ffe9..760bf202d1 100644
--- a/src/components/entity/ha-statistic-picker.ts
+++ b/src/components/entity/ha-statistic-picker.ts
@@ -13,7 +13,7 @@ import { computeAreaName } from "../../common/entity/compute_area_name";
 import { computeDeviceName } from "../../common/entity/compute_device_name";
 import { computeEntityName } from "../../common/entity/compute_entity_name";
 import { computeStateName } from "../../common/entity/compute_state_name";
-import { getEntityContext } from "../../common/entity/get_entity_context";
+import { getEntityContext } from "../../common/entity/context/get_entity_context";
 import { caseInsensitiveStringCompare } from "../../common/string/compare";
 import { computeRTL } from "../../common/util/compute_rtl";
 import { domainToName } from "../../data/integration";
diff --git a/src/components/ha-areas-display-editor.ts b/src/components/ha-areas-display-editor.ts
index 4215b08c0c..9a63729ccf 100644
--- a/src/components/ha-areas-display-editor.ts
+++ b/src/components/ha-areas-display-editor.ts
@@ -44,7 +44,7 @@ export class HaAreasDisplayEditor extends LitElement {
     );
 
     const items: DisplayItem[] = areas.map((area) => {
-      const { floor } = getAreaContext(area.area_id, this.hass!);
+      const { floor } = getAreaContext(area, this.hass!);
       return {
         value: area.area_id,
         label: area.name,
diff --git a/src/dialogs/more-info/ha-more-info-dialog.ts b/src/dialogs/more-info/ha-more-info-dialog.ts
index c8a1c34437..a15d2d7765 100644
--- a/src/dialogs/more-info/ha-more-info-dialog.ts
+++ b/src/dialogs/more-info/ha-more-info-dialog.ts
@@ -25,6 +25,10 @@ import {
   computeEntityEntryName,
   computeEntityName,
 } from "../../common/entity/compute_entity_name";
+import {
+  getEntityContext,
+  getEntityEntryContext,
+} from "../../common/entity/context/get_entity_context";
 import { shouldHandleRequestSelectedEvent } from "../../common/mwc/handle-request-selected-event";
 import { navigate } from "../../common/navigate";
 import "../../components/ha-button-menu";
@@ -58,10 +62,6 @@ import "./ha-more-info-history-and-logbook";
 import "./ha-more-info-info";
 import "./ha-more-info-settings";
 import "./more-info-content";
-import {
-  getEntityContext,
-  getEntityEntryContext,
-} from "../../common/entity/get_entity_context";
 
 export interface MoreInfoDialogParams {
   entityId: string | null;
diff --git a/src/panels/config/dashboard/ha-config-updates.ts b/src/panels/config/dashboard/ha-config-updates.ts
index 539655c3f3..adedf303c4 100644
--- a/src/panels/config/dashboard/ha-config-updates.ts
+++ b/src/panels/config/dashboard/ha-config-updates.ts
@@ -6,7 +6,7 @@ import { ifDefined } from "lit/directives/if-defined";
 import memoizeOne from "memoize-one";
 import { fireEvent } from "../../../common/dom/fire_event";
 import { computeDeviceNameDisplay } from "../../../common/entity/compute_device_name";
-import { getDeviceContext } from "../../../common/entity/get_device_context";
+import { getDeviceContext } from "../../../common/entity/context/get_device_context";
 import "../../../components/entity/state-badge";
 import "../../../components/ha-alert";
 import "../../../components/ha-icon-next";
diff --git a/test/common/entity/context/context-mock.ts b/test/common/entity/context/context-mock.ts
new file mode 100644
index 0000000000..6fa8e5dce9
--- /dev/null
+++ b/test/common/entity/context/context-mock.ts
@@ -0,0 +1,85 @@
+import type { HassEntity } from "home-assistant-js-websocket";
+import type { EntityRegistryDisplayEntry } from "../../../../src/data/entity_registry";
+import type { DeviceRegistryEntry } from "../../../../src/data/device_registry";
+import type { AreaRegistryEntry } from "../../../../src/data/area_registry";
+import type { FloorRegistryEntry } from "../../../../src/data/floor_registry";
+
+export const mockStateObj = (partial: Partial<HassEntity>): HassEntity => ({
+  entity_id: "",
+  attributes: {},
+  state: "on",
+  last_changed: "",
+  last_updated: "",
+  context: {
+    id: "",
+    user_id: null,
+    parent_id: null,
+  },
+  ...partial,
+});
+
+export const mockEntity = (
+  partial: Partial<EntityRegistryDisplayEntry>
+): EntityRegistryDisplayEntry => ({
+  entity_id: "",
+  labels: [],
+  ...partial,
+});
+
+export const mockDevice = (
+  partial: Partial<DeviceRegistryEntry>
+): DeviceRegistryEntry => ({
+  id: "",
+  config_entries: [],
+  config_entries_subentries: {},
+  connections: [],
+  identifiers: [],
+  manufacturer: null,
+  model: null,
+  model_id: null,
+  name: null,
+  labels: [],
+  sw_version: null,
+  hw_version: null,
+  serial_number: null,
+  via_device_id: null,
+  area_id: null,
+  name_by_user: null,
+  entry_type: null,
+  disabled_by: null,
+  configuration_url: null,
+  primary_config_entry: null,
+  created_at: 0,
+  modified_at: 0,
+  ...partial,
+});
+
+export const mockArea = (
+  partial: Partial<AreaRegistryEntry>
+): AreaRegistryEntry => ({
+  aliases: [],
+  area_id: "",
+  name: "",
+  floor_id: null,
+  created_at: 0,
+  modified_at: 0,
+  humidity_entity_id: null,
+  temperature_entity_id: null,
+  icon: null,
+  labels: [],
+  picture: null,
+  ...partial,
+});
+
+export const mockFloor = (
+  partial: Partial<FloorRegistryEntry>
+): FloorRegistryEntry => ({
+  aliases: [],
+  floor_id: "",
+  name: "",
+  created_at: 0,
+  modified_at: 0,
+  icon: null,
+  level: 0,
+  ...partial,
+});
diff --git a/test/common/entity/context/get_area_context.test.ts b/test/common/entity/context/get_area_context.test.ts
index d3e26631d3..a35ea5b9e5 100644
--- a/test/common/entity/context/get_area_context.test.ts
+++ b/test/common/entity/context/get_area_context.test.ts
@@ -1,53 +1,53 @@
 import { describe, it, expect } from "vitest";
 import { getAreaContext } from "../../../../src/common/entity/context/get_area_context";
 import type { HomeAssistant } from "../../../../src/types";
+import { mockArea, mockFloor } from "./context-mock";
 
 describe("getAreaContext", () => {
-  it("should return null values when the area does not exist", () => {
-    const hass = {
-      areas: {},
-      floors: {},
-    } as unknown as HomeAssistant;
-
-    const result = getAreaContext("nonexistent.area", hass);
-
-    expect(result).toEqual({
-      area: null,
-      floor: null,
-    });
-  });
-
   it("should return the correct context when the area exists without a floor", () => {
+    const area = mockArea({
+      area_id: "area_1",
+    });
+
     const hass = {
       areas: {
-        area_1: { id: "area_1" },
+        area_1: area,
       },
       floors: {},
     } as unknown as HomeAssistant;
 
-    const result = getAreaContext("area_1", hass);
+    const result = getAreaContext(area, hass);
 
     expect(result).toEqual({
-      area: { id: "area_1" },
+      area,
       floor: null,
     });
   });
 
   it("should return the correct context when the area exists with a floor", () => {
+    const area = mockArea({
+      area_id: "area_2",
+      floor_id: "floor_1",
+    });
+
+    const floor = mockFloor({
+      floor_id: "floor_1",
+    });
+
     const hass = {
       areas: {
-        area_2: { id: "area_2", floor_id: "floor_1" },
+        area_2: area,
       },
       floors: {
-        floor_1: { id: "floor_1" },
+        floor_1: floor,
       },
     } as unknown as HomeAssistant;
 
-    const result = getAreaContext("area_2", hass);
+    const result = getAreaContext(area, hass);
 
     expect(result).toEqual({
-      area: { id: "area_2", floor_id: "floor_1" },
-      floor: { id: "floor_1" },
+      area,
+      floor,
     });
   });
 });
diff --git a/test/common/entity/context/get_device_context.test.ts b/test/common/entity/context/get_device_context.test.ts
new file mode 100644
index 0000000000..83c66edaba
--- /dev/null
+++ b/test/common/entity/context/get_device_context.test.ts
@@ -0,0 +1,93 @@
+import { describe, expect, it } from "vitest";
+import { getDeviceContext } from "../../../../src/common/entity/context/get_device_context";
+import type { HomeAssistant } from "../../../../src/types";
+import { mockArea, mockDevice, mockFloor } from "./context-mock";
+
+describe("getDeviceContext", () => {
+  it("should return the correct context when the device exists without area", () => {
+    const device = mockDevice({
+      id: "device_1",
+    });
+
+    const hass = {
+      devices: {
+        device_1: device,
+      },
+      areas: {},
+      floors: {},
+    } as unknown as HomeAssistant;
+
+    const result = getDeviceContext(device, hass);
+
+    expect(result).toEqual({
+      device,
+      area: null,
+      floor: null,
+    });
+  });
+
+  it("should return the correct context when the device exists with area but no floor", () => {
+    const device = mockDevice({
+      id: "device_2",
+      area_id: "area_1",
+    });
+
+    const area = mockArea({
+      area_id: "area_1",
+    });
+
+    const hass = {
+      devices: {
+        device_2: device,
+      },
+      areas: {
+        area_1: area,
+      },
+      floors: {},
+    } as unknown as HomeAssistant;
+
+    const result = getDeviceContext(device, hass);
+
+    expect(result).toEqual({
+      device,
+      area,
+      floor: null,
+    });
+  });
+
+  it("should return the correct context when the device exists with area and floor", () => {
+    const device = mockDevice({
+      id: "device_3",
+      area_id: "area_2",
+    });
+
+    const area = mockArea({
+      area_id: "area_2",
+      floor_id: "floor_1",
+    });
+
+    const floor = mockFloor({
+      floor_id: "floor_1",
+    });
+
+    const hass = {
+      devices: {
+        device_3: device,
+      },
+      areas: {
+        area_2: area,
+      },
+      floors: {
+        floor_1: floor,
+      },
+    } as unknown as HomeAssistant;
+
+    const result = getDeviceContext(device, hass);
+
+    expect(result).toEqual({
+      device,
+      area,
+      floor,
+    });
+  });
+});
diff --git a/test/common/entity/context/get_entity_context.test.ts b/test/common/entity/context/get_entity_context.test.ts
index af0077225c..82a74b5e76 100644
--- a/test/common/entity/context/get_entity_context.test.ts
+++ b/test/common/entity/context/get_entity_context.test.ts
@@ -1,40 +1,35 @@
-import { describe, it, expect } from "vitest";
-import type { HomeAssistant } from "../../../../src/types";
+import { describe, expect, it } from "vitest";
 import { getEntityContext } from "../../../../src/common/entity/context/get_entity_context";
+import type { HomeAssistant } from "../../../../src/types";
+import {
+  mockArea,
+  mockDevice,
+  mockEntity,
+  mockFloor,
+  mockStateObj,
+} from "./context-mock";
 
 describe("getEntityContext", () => {
-  it("should return null values when the entity does not exist", () => {
-    const hass = {
-      entities: {},
-      devices: {},
-      areas: {},
-      floors: {},
-    } as unknown as HomeAssistant;
-
-    const result = getEntityContext("nonexistent.entity", hass);
-
-    expect(result).toEqual({
-      entity: null,
-      device: null,
-      area: null,
-      floor: null,
-    });
-  });
-
   it("should return the correct context when the entity exists without device or area", () => {
+    const entity = mockEntity({
+      entity_id: "light.living_room",
+    });
+    const stateObj = mockStateObj({
+      entity_id: "light.living_room",
+    });
     const hass = {
       entities: {
-        "light.living_room": { entity_id: "light.living_room" },
+        "light.living_room": entity,
       },
       devices: {},
       areas: {},
       floors: {},
     } as unknown as HomeAssistant;
 
-    const result = getEntityContext("light.living_room", hass);
+    const result = getEntityContext(stateObj, hass);
 
     expect(result).toEqual({
-      entity: { entity_id: "light.living_room" },
+      entity,
       device: null,
       area: null,
       floor: null,
@@ -42,76 +37,113 @@ describe("getEntityContext", () => {
   });
 
   it("should return the correct context when the entity has a device and area", () => {
+    const entity = mockEntity({
+      entity_id: "light.living_room",
+      device_id: "device_1",
+    });
+    const device = mockDevice({
+      id: "device_1",
+      area_id: "area_1",
+    });
+    const area = mockArea({
+      area_id: "area_1",
+      floor_id: "floor_1",
+    });
+    const floor = mockFloor({
+      floor_id: "floor_1",
+    });
+    const stateObj = mockStateObj({
+      entity_id: "light.living_room",
+    });
+
     const hass = {
       entities: {
-        "light.living_room": {
-          entity_id: "light.living_room",
-          device_id: "device_1",
-        },
+        "light.living_room": entity,
       },
       devices: {
-        device_1: { id: "device_1", area_id: "area_1" },
+        device_1: device,
       },
       areas: {
-        area_1: { id: "area_1", floor_id: "floor_1" },
+        area_1: area,
       },
       floors: {
-        floor_1: { id: "floor_1" },
+        floor_1: floor,
       },
     } as unknown as HomeAssistant;
 
-    const result = getEntityContext("light.living_room", hass);
+    const result = getEntityContext(stateObj, hass);
 
     expect(result).toEqual({
-      entity: { entity_id: "light.living_room", device_id: "device_1" },
-      device: { id: "device_1", area_id: "area_1" },
-      area: { id: "area_1", floor_id: "floor_1" },
-      floor: { id: "floor_1" },
+      entity,
+      device,
+      area,
+      floor,
     });
   });
 
   it("should return the correct context when the entity has an area but no device", () => {
+    const entity = mockEntity({
+      entity_id: "sensor.kitchen",
+      area_id: "area_2",
+    });
+    const area = mockArea({ area_id: "area_2", floor_id: "floor_2" });
+    const floor = mockFloor({ floor_id: "floor_2" });
+    const stateObj = mockStateObj({
+      entity_id: "sensor.kitchen",
+    });
+
     const hass = {
       entities: {
-        "sensor.kitchen": { entity_id: "sensor.kitchen", area_id: "area_2" },
+        "sensor.kitchen": entity,
       },
       devices: {},
       areas: {
-        area_2: { id: "area_2", floor_id: "floor_2" },
+        area_2: area,
       },
       floors: {
-        floor_2: { id: "floor_2" },
+        floor_2: floor,
       },
     } as unknown as HomeAssistant;
 
-    const result = getEntityContext("sensor.kitchen", hass);
+    const result = getEntityContext(stateObj, hass);
 
     expect(result).toEqual({
-      entity: { entity_id: "sensor.kitchen", area_id: "area_2" },
+      entity,
       device: null,
-      area: { id: "area_2", floor_id: "floor_2" },
-      floor: { id: "floor_2" },
+      area,
+      floor,
     });
   });
 
   it("should return null for floor if area does not have a floor_id", () => {
+    const entity = mockEntity({
+      entity_id: "sensor.bedroom",
+      area_id: "area_3",
+    });
+    const area = mockArea({
+      area_id: "area_3",
+    });
+    const stateObj = mockStateObj({
+      entity_id: "sensor.bedroom",
+    });
+
     const hass = {
       entities: {
-        "sensor.bedroom": { entity_id: "sensor.bedroom", area_id: "area_3" },
+        "sensor.bedroom": entity,
       },
       devices: {},
       areas: {
-        area_3: { id: "area_3" },
+        area_3: area,
       },
       floors: {},
     } as unknown as HomeAssistant;
 
-    const result = getEntityContext("sensor.bedroom", hass);
+    const result = getEntityContext(stateObj, hass);
 
     expect(result).toEqual({
-      entity: { entity_id: "sensor.bedroom", area_id: "area_3" },
+      entity,
       device: null,
-      area: { id: "area_3" },
+      area,
       floor: null,
     });
   });