diff --git a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/constants.ts b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/constants.ts
index 4dff70518c1157..910016d3871630 100644
--- a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/constants.ts
+++ b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/constants.ts
@@ -39,3 +39,58 @@ export const DELETE_PHASE_POLICY: PolicyFromES = {
   },
   name: POLICY_NAME,
 };
+
+export const POLICY_WITH_NODE_ATTR_AND_OFF_ALLOCATION: PolicyFromES = {
+  version: 1,
+  modified_date: Date.now().toString(),
+  policy: {
+    phases: {
+      hot: {
+        min_age: '0ms',
+        actions: {
+          rollover: {
+            max_size: '50gb',
+          },
+        },
+      },
+      warm: {
+        actions: {
+          allocate: {
+            require: {},
+            include: { test: '123' },
+            exclude: {},
+          },
+        },
+      },
+      cold: {
+        actions: {
+          migrate: { enabled: false },
+        },
+      },
+    },
+    name: POLICY_NAME,
+  },
+  name: POLICY_NAME,
+};
+
+export const POLICY_WITH_NODE_ROLE_ALLOCATION: PolicyFromES = {
+  version: 1,
+  modified_date: Date.now().toString(),
+  policy: {
+    phases: {
+      hot: {
+        min_age: '0ms',
+        actions: {
+          rollover: {
+            max_size: '50gb',
+          },
+        },
+      },
+      warm: {
+        actions: {},
+      },
+    },
+    name: POLICY_NAME,
+  },
+  name: POLICY_NAME,
+};
diff --git a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.helpers.tsx b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.helpers.tsx
index 6365bb8caa9631..dc076dd36117bf 100644
--- a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.helpers.tsx
+++ b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.helpers.tsx
@@ -45,7 +45,7 @@ const testBedConfig: TestBedConfig = {
 
 const initTestBed = registerTestBed(EditPolicy, testBedConfig);
 
-export interface EditPolicyTestBed extends TestBed<TestSubjects> {
+export interface EditPolicyTestBed extends TestBed<TestSubjects | string> {
   actions: {
     setWaitForSnapshotPolicy: (snapshotPolicyName: string) => void;
     savePolicy: () => void;
diff --git a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.test.ts b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.test.ts
index b465afb8b5d9f4..5bcb48ecfbe467 100644
--- a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.test.ts
+++ b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/edit_policy/edit_policy.test.ts
@@ -10,7 +10,13 @@ import { setupEnvironment } from '../helpers/setup_environment';
 import { EditPolicyTestBed, setup } from './edit_policy.helpers';
 
 import { API_BASE_PATH } from '../../../common/constants';
-import { DELETE_PHASE_POLICY, NEW_SNAPSHOT_POLICY_NAME, SNAPSHOT_POLICY_NAME } from './constants';
+import {
+  DELETE_PHASE_POLICY,
+  NEW_SNAPSHOT_POLICY_NAME,
+  SNAPSHOT_POLICY_NAME,
+  POLICY_WITH_NODE_ATTR_AND_OFF_ALLOCATION,
+  POLICY_WITH_NODE_ROLE_ALLOCATION,
+} from './constants';
 
 window.scrollTo = jest.fn();
 
@@ -134,4 +140,57 @@ describe('<EditPolicy />', () => {
       expect(testBed.find('policiesErrorCallout').exists()).toBeTruthy();
     });
   });
+
+  describe('data allocation', () => {
+    describe('node roles', () => {
+      beforeEach(async () => {
+        httpRequestsMockHelpers.setLoadPolicies([POLICY_WITH_NODE_ROLE_ALLOCATION]);
+        httpRequestsMockHelpers.setListNodes({
+          isUsingDeprecatedDataRoleConfig: false,
+          nodesByAttributes: { test: ['123'] },
+          nodesByRoles: { data: ['123'] },
+        });
+
+        await act(async () => {
+          testBed = await setup();
+        });
+
+        const { component } = testBed;
+        component.update();
+      });
+      test('showing "default" type', () => {
+        const { find } = testBed;
+        expect(find('warm-dataTierAllocationControls.dataTierSelect').text()).toContain(
+          'recommended'
+        );
+        expect(find('warm-dataTierAllocationControls.dataTierSelect').text()).not.toContain(
+          'Custom'
+        );
+        expect(find('warm-dataTierAllocationControls.dataTierSelect').text()).not.toContain('Off');
+      });
+    });
+    describe('node attr and none', () => {
+      beforeEach(async () => {
+        httpRequestsMockHelpers.setLoadPolicies([POLICY_WITH_NODE_ATTR_AND_OFF_ALLOCATION]);
+        httpRequestsMockHelpers.setListNodes({
+          isUsingDeprecatedDataRoleConfig: false,
+          nodesByAttributes: { test: ['123'] },
+          nodesByRoles: { data: ['123'] },
+        });
+
+        await act(async () => {
+          testBed = await setup();
+        });
+
+        const { component } = testBed;
+        component.update();
+      });
+
+      test('showing "custom" and "off" types', () => {
+        const { find } = testBed;
+        expect(find('warm-dataTierAllocationControls.dataTierSelect').text()).toContain('Custom');
+        expect(find('cold-dataTierAllocationControls.dataTierSelect').text()).toContain('Off');
+      });
+    });
+  });
 });
diff --git a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/helpers/http_requests.ts b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/helpers/http_requests.ts
index 04f58f93939ca3..c7a493ce80d96b 100644
--- a/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/helpers/http_requests.ts
+++ b/x-pack/plugins/index_lifecycle_management/__jest__/client_integration/helpers/http_requests.ts
@@ -6,6 +6,7 @@
 
 import { fakeServer, SinonFakeServer } from 'sinon';
 import { API_BASE_PATH } from '../../../common/constants';
+import { ListNodesRouteResponse } from '../../../common/types';
 
 export const init = () => {
   const server = fakeServer.create();
@@ -38,8 +39,17 @@ const registerHttpRequestMockHelpers = (server: SinonFakeServer) => {
     ]);
   };
 
+  const setListNodes = (body: ListNodesRouteResponse) => {
+    server.respondWith('GET', `${API_BASE_PATH}/nodes/list`, [
+      200,
+      { 'Content-Type': 'application/json' },
+      JSON.stringify(body),
+    ]);
+  };
+
   return {
     setLoadPolicies,
     setLoadSnapshotPolicies,
+    setListNodes,
   };
 };