diff --git a/demo/src/ha-demo.ts b/demo/src/ha-demo.ts
index 958e14e0e3..91edb70aa5 100644
--- a/demo/src/ha-demo.ts
+++ b/demo/src/ha-demo.ts
@@ -21,8 +21,9 @@ import { mockSystemLog } from "./stubs/system_log";
 import { mockTemplate } from "./stubs/template";
 import { mockTranslations } from "./stubs/translations";
 import { mockEnergy } from "./stubs/energy";
-import { mockConfig } from "./stubs/config";
 import { energyEntities } from "./stubs/entities";
+import { mockConfigEntries } from "./stubs/config_entries";
+import { mockEntityRegistry } from "./stubs/entity_registry";
 
 class HaDemo extends HomeAssistantAppEl {
   protected async _initializeHass() {
@@ -51,8 +52,34 @@ class HaDemo extends HomeAssistantAppEl {
     mockMediaPlayer(hass);
     mockFrontend(hass);
     mockEnergy(hass);
-    mockConfig(hass);
     mockPersistentNotification(hass);
+    mockConfigEntries(hass);
+    mockEntityRegistry(hass, [
+      {
+        config_entry_id: "co2signal",
+        device_id: "co2signal",
+        area_id: null,
+        disabled_by: null,
+        entity_id: "sensor.co2_intensity",
+        name: null,
+        icon: null,
+        platform: "co2signal",
+        hidden_by: null,
+        entity_category: null,
+      },
+      {
+        config_entry_id: "co2signal",
+        device_id: "co2signal",
+        area_id: null,
+        disabled_by: null,
+        entity_id: "sensor.grid_fossil_fuel_percentage",
+        name: null,
+        icon: null,
+        platform: "co2signal",
+        hidden_by: null,
+        entity_category: null,
+      },
+    ]);
 
     hass.addEntities(energyEntities());
 
diff --git a/demo/src/stubs/config.ts b/demo/src/stubs/config.ts
deleted file mode 100644
index f77c8d3b09..0000000000
--- a/demo/src/stubs/config.ts
+++ /dev/null
@@ -1,41 +0,0 @@
-import { MockHomeAssistant } from "../../../src/fake_data/provide_hass";
-
-export const mockConfig = (hass: MockHomeAssistant) => {
-  hass.mockAPI("config/config_entries/entry?domain=co2signal", () => [
-    {
-      entry_id: "co2signal",
-      domain: "co2signal",
-      title: "CO2 Signal",
-      source: "user",
-      state: "loaded",
-      supports_options: false,
-      supports_unload: true,
-      pref_disable_new_entities: false,
-      pref_disable_polling: false,
-      disabled_by: null,
-      reason: null,
-    },
-  ]);
-  hass.mockWS("config/entity_registry/list", () => [
-    {
-      config_entry_id: "co2signal",
-      device_id: "co2signal",
-      area_id: null,
-      disabled_by: null,
-      entity_id: "sensor.co2_intensity",
-      name: null,
-      icon: null,
-      platform: "co2signal",
-    },
-    {
-      config_entry_id: "co2signal",
-      device_id: "co2signal",
-      area_id: null,
-      disabled_by: null,
-      entity_id: "sensor.grid_fossil_fuel_percentage",
-      name: null,
-      icon: null,
-      platform: "co2signal",
-    },
-  ]);
-};
diff --git a/demo/src/stubs/config_entries.ts b/demo/src/stubs/config_entries.ts
new file mode 100644
index 0000000000..f81e11bc39
--- /dev/null
+++ b/demo/src/stubs/config_entries.ts
@@ -0,0 +1,20 @@
+import { MockHomeAssistant } from "../../../src/fake_data/provide_hass";
+
+export const mockConfigEntries = (hass: MockHomeAssistant) => {
+  hass.mockWS("config_entries/get", () => [
+    {
+      entry_id: "co2signal",
+      domain: "co2signal",
+      title: "CO2 Signal",
+      source: "user",
+      state: "loaded",
+      supports_options: false,
+      supports_remove_device: false,
+      supports_unload: true,
+      pref_disable_new_entities: false,
+      pref_disable_polling: false,
+      disabled_by: null,
+      reason: null,
+    },
+  ]);
+};
diff --git a/demo/src/stubs/entity_registry.ts b/demo/src/stubs/entity_registry.ts
index 8f548629e7..422702b646 100644
--- a/demo/src/stubs/entity_registry.ts
+++ b/demo/src/stubs/entity_registry.ts
@@ -4,4 +4,6 @@ import type { MockHomeAssistant } from "../../../src/fake_data/provide_hass";
 export const mockEntityRegistry = (
   hass: MockHomeAssistant,
   data: EntityRegistryEntry[] = []
-) => hass.mockWS("config/entity_registry/list", () => data);
+) => {
+  hass.mockWS("config/entity_registry/list", () => data);
+};