diff --git a/src/Umbraco.Web.UI.Client/src/mocks/data/entity.data.ts b/src/Umbraco.Web.UI.Client/src/mocks/data/entity.data.ts
index a8ca0acd35..9b1bf730ac 100644
--- a/src/Umbraco.Web.UI.Client/src/mocks/data/entity.data.ts
+++ b/src/Umbraco.Web.UI.Client/src/mocks/data/entity.data.ts
@@ -1,8 +1,10 @@
+import { UmbData } from './data';
+
 export interface Entity {
 	id: number;
 	key: string;
 	name: string;
-	icon: string; // TODO: Should this be here?
+	icon?: string; // TODO: Should this be here?
 	type: string;
 	hasChildren: boolean; // TODO: Should this be here?
 	parentKey: string;
@@ -11,38 +13,39 @@ export interface Entity {
 export const data: Array<Entity> = [
 	{
 		id: 1,
-		key: '74e4008a-ea4f-4793-b924-15e02fd380d1',
-		parentKey: '',
-		name: 'Document 1',
-		type: 'document',
-		icon: 'document',
+		key: '865a11f9-d140-4f21-8dfe-2caafc65a971',
+		type: 'member',
+		parentKey: '24fcd88a-d1bb-423b-b794-8a94dcddcb6a',
+		name: 'Member 1',
 		hasChildren: false,
 	},
 	{
 		id: 2,
-		key: '74e4008a-ea4f-4793-b924-15e02fd380d2',
-		parentKey: '',
-		name: 'Document 2',
-		type: 'document',
-		icon: 'favorite',
-		hasChildren: false,
+		key: '06c6919c-6fa7-4aa5-8214-0582c721c472',
+		type: 'member',
+		parentKey: '24fcd88a-d1bb-423b-b794-8a94dcddcb6a',
+		name: 'Member 2',
+		hasChildren: true,
 	},
 	{
-		id: 3,
-		key: 'cdd30288-2d1c-41b4-89a9-61647b4a10d5',
-		parentKey: '',
-		name: 'Document 3',
-		type: 'document',
-		icon: 'document',
-		hasChildren: false,
-	},
-	{
-		id: 2001,
-		key: 'f2f81a40-c989-4b6b-84e2-057cecd3adc1',
-		parentKey: '',
-		name: 'Media 1',
-		type: 'media',
-		icon: 'picture',
+		id: 2,
+		key: '725a26c4-158d-4dc0-8aaa-b64473b11aa8',
+		type: 'member',
+		parentKey: '06c6919c-6fa7-4aa5-8214-0582c721c472',
+		name: 'Member 3',
 		hasChildren: false,
 	},
 ];
+
+// Temp mocked database
+class UmbEntityData extends UmbData<Entity> {
+	constructor() {
+		super(data);
+	}
+
+	getChildren(key: string) {
+		return data.filter((item) => item.parentKey === key);
+	}
+}
+
+export const umbEntityData = new UmbEntityData();
diff --git a/src/Umbraco.Web.UI.Client/src/mocks/domains/tree.handlers.ts b/src/Umbraco.Web.UI.Client/src/mocks/domains/tree.handlers.ts
index 42be22d732..532ea5df65 100644
--- a/src/Umbraco.Web.UI.Client/src/mocks/domains/tree.handlers.ts
+++ b/src/Umbraco.Web.UI.Client/src/mocks/domains/tree.handlers.ts
@@ -1,128 +1,32 @@
 import { rest } from 'msw';
-import { Entity } from '../data/entity.data';
+import { Entity, umbEntityData } from '../data/entity.data';
 
 // TODO: add schema
 export const handlers = [
-	rest.get('/umbraco/backoffice/trees/members/:id', (req, res, ctx) => {
+	rest.get('/umbraco/backoffice/trees/members/:key', (req, res, ctx) => {
 		console.warn('Please move to schema');
-		const id = req.params.id as string;
-		if (!id) return;
+		const key = req.params.key as string;
+		if (!key) return;
 
-		//const int = parseInt(id);
-		//const document = umbNodeData.getById(int);
-
-		const items = [
-			{
-				id: '1',
-				key: '865a11f9-d140-4f21-8dfe-2caafc65a971',
-				parentKey: '24fcd88a-d1bb-423b-b794-8a94dcddcb6a',
-				name: 'Member 1',
-				hasChildren: false,
-			},
-			{
-				id: '2',
-				key: '06c6919c-6fa7-4aa5-8214-0582c721c472',
-				parentKey: '24fcd88a-d1bb-423b-b794-8a94dcddcb6a',
-				name: 'Member 2',
-				hasChildren: true,
-			},
-			{
-				id: '2',
-				key: '725a26c4-158d-4dc0-8aaa-b64473b11aa8',
-				parentKey: '06c6919c-6fa7-4aa5-8214-0582c721c472',
-				name: 'Member 3',
-				hasChildren: false,
-			},
-		];
-
-		const test = items.filter((item) => item.parentKey === id);
-
-		return res(ctx.status(200), ctx.json(test));
+		const entities = umbEntityData.getChildren(key);
+		return res(ctx.status(200), ctx.json(entities));
 	}),
 
-	rest.get('/umbraco/backoffice/trees/member-groups/:id', (req, res, ctx) => {
+	rest.get('/umbraco/backoffice/trees/member-groups/:key', (req, res, ctx) => {
 		console.warn('Please move to schema');
-		const id = req.params.id as string;
-		if (!id) return;
+		const key = req.params.key as string;
+		if (!key) return;
 
-		//const int = parseInt(id);
-		//const document = umbNodeData.getById(int);
-
-		const items = [
-			{
-				id: '1',
-				key: 'fcb3f468-97de-469d-8090-d14d068ad968',
-				name: 'Group 1',
-				hasChildren: false,
-			},
-			{
-				id: '2',
-				key: 'd99bfca4-551d-427d-a842-d47b756b8977',
-				name: 'Group 2',
-				hasChildren: false,
-			},
-			{
-				id: '3',
-				key: 'b7b7ec7c-1c4a-4a78-9a6a-0652e891972a',
-				name: 'Group 3',
-				hasChildren: true,
-			},
-		];
-
-		return res(ctx.status(200), ctx.json(items));
+		const entities = umbEntityData.getChildren(key);
+		return res(ctx.status(200), ctx.json(entities));
 	}),
 
-	rest.get('/umbraco/backoffice/trees/data-types/:id', (req, res, ctx) => {
+	rest.get('/umbraco/backoffice/trees/data-types/:key', (req, res, ctx) => {
 		console.warn('Please move to schema');
-		const id = req.params.id as string;
-		if (!id) return;
+		const key = req.params.key as string;
+		if (!key) return;
 
-		//const int = parseInt(id);
-		//const document = umbNodeData.getById(int);
-
-		const items: Array<Entity> = [
-			{
-				id: 1245,
-				key: 'dt-1',
-				type: 'data-type',
-				name: 'Text',
-				hasChildren: false,
-				icon: 'favorite',
-			},
-			{
-				id: 1244,
-				key: 'dt-2',
-				type: 'data-type',
-				name: 'Textarea',
-				hasChildren: false,
-				icon: 'favorite',
-			},
-			{
-				id: 1246,
-				key: 'dt-3',
-				type: 'data-type',
-				name: 'My JS Property Editor',
-				hasChildren: false,
-				icon: 'favorite',
-			},
-			{
-				id: 1247,
-				key: 'dt-4',
-				type: 'data-type',
-				name: 'Context Example',
-				hasChildren: false,
-				icon: 'favorite',
-			},
-			{
-				id: 1248,
-				key: 'dt-5',
-				type: 'data-type',
-				name: 'Content Picker (DataType)',
-				hasChildren: false,
-				icon: 'favorite',
-			},
-		];
-
-		return res(ctx.status(200), ctx.json(items));
+		const entities = umbEntityData.getChildren(key);
+		return res(ctx.status(200), ctx.json(entities));
 	}),
 ];