@@ -109,6 +109,32 @@ TEST_F(PoolResidencyTestFixture,
109109 umfPoolFree (pool, ptr);
110110}
111111
112+ TEST_F (PoolResidencyTestFixture,
113+ existingAllocationsShouldBeEvictedFromRemovedDevice) {
114+ initializeMemoryPool (l0mock.initializeMemoryProviderWithResidentDevices (
115+ OUR_DEVICE, {DEVICE_2, DEVICE_3}));
116+
117+ EXPECT_CALL (l0mock, zeMemAllocDevice (CONTEXT, _, _, _, OUR_DEVICE, _))
118+ .WillOnce (
119+ DoAll (SetArgPointee<5 >(POINTER_0), Return (ZE_RESULT_SUCCESS)));
120+ EXPECT_CALL (l0mock, zeContextMakeMemoryResident (CONTEXT, DEVICE_2, _, _))
121+ .WillOnce (Return (ZE_RESULT_SUCCESS));
122+ EXPECT_CALL (l0mock, zeContextMakeMemoryResident (CONTEXT, DEVICE_3, _, _))
123+ .WillOnce (Return (ZE_RESULT_SUCCESS));
124+
125+ void *ptr = umfPoolMalloc (pool, 123 );
126+ EXPECT_EQ (ptr, POINTER_0);
127+
128+ EXPECT_CALL (l0mock, zeContextEvictMemory (CONTEXT, DEVICE_2, _, _))
129+ .WillOnce (Return (ZE_RESULT_SUCCESS));
130+
131+ umf_memory_provider_handle_t provider = nullptr ;
132+ EXPECT_EQ (umfPoolGetMemoryProvider (pool, &provider), UMF_RESULT_SUCCESS);
133+ umfLevelZeroMemoryProviderResidentDeviceChange (provider, DEVICE_2, false );
134+
135+ umfPoolFree (pool, ptr);
136+ }
137+
112138TEST_F (PoolResidencyTestFixture,
113139 allocationShouldNotBeMadeResidentOnRemovedDevice) {
114140 initializeMemoryPool (l0mock.initializeMemoryProviderWithResidentDevices (
0 commit comments