diff --git a/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java b/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java
index 3fd0e072809ff..f8cfdf131a28d 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java
@@ -17,9 +17,12 @@
 package com.android.server.am;
 
 import static android.app.ActivityManager.PROCESS_STATE_BOUND_FOREGROUND_SERVICE;
+
 import static com.android.server.am.ActivityManagerService.Injector;
 import static com.android.server.am.CachedAppOptimizer.compactActionIntToAction;
+
 import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
 
@@ -34,27 +37,30 @@ import android.os.Process;
 import android.platform.test.annotations.Presubmit;
 import android.provider.DeviceConfig;
 import android.text.TextUtils;
+
 import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.dx.mockito.inline.extended.StaticMockitoSessionBuilder;
 import com.android.modules.utils.testing.TestableDeviceConfig;
+import com.android.server.ExtendedMockitoTestCase;
 import com.android.server.LocalServices;
 import com.android.server.ServiceThread;
 import com.android.server.appop.AppOpsService;
 import com.android.server.wm.ActivityTaskManagerService;
+
+import org.junit.After;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+
 import java.io.File;
 import java.io.IOException;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import org.junit.After;
-import org.junit.Assume;
-import org.junit.Before;
-import org.junit.Ignore;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnitRunner;
 
 /**
  * Tests for {@link CachedAppOptimizer}.
@@ -63,10 +69,7 @@ import org.mockito.junit.MockitoJUnitRunner;
  * atest FrameworksMockingServicesTests:CachedAppOptimizerTest
  */
 @Presubmit
-@RunWith(MockitoJUnitRunner.class)
-@Ignore("TODO(b/226641572): this test is broken and it cannot use ExtendedMockitoTestCase as it "
-        + "uses TestableDeviceConfigRule, which creates its own mockito session")
-public final class CachedAppOptimizerTest {
+public final class CachedAppOptimizerTest extends ExtendedMockitoTestCase {
 
     private ServiceThread mThread;
 
@@ -84,16 +87,21 @@ public final class CachedAppOptimizerTest {
     @Mock
     private PackageManagerInternal mPackageManagerInt;
 
-    @Rule
-    public final TestableDeviceConfig.TestableDeviceConfigRule
-            mDeviceConfigRule = new TestableDeviceConfig.TestableDeviceConfigRule();
+    private final TestableDeviceConfig mDeviceConfig = new TestableDeviceConfig();
+
     @Rule
     public final ApplicationExitInfoTest.ServiceThreadRule
             mServiceThreadRule = new ApplicationExitInfoTest.ServiceThreadRule();
 
+    @Override
+    protected void initializeSession(StaticMockitoSessionBuilder builder) {
+        mDeviceConfig.setUpMockedClasses(builder);
+    }
+
     @Before
     public void setUp() {
         System.loadLibrary("mockingservicestestjni");
+        mDeviceConfig.setUpMockBehaviors();
         mHandlerThread = new HandlerThread("");
         mHandlerThread.start();
         mHandler = new Handler(mHandlerThread.getLooper());
@@ -124,6 +132,7 @@ public final class CachedAppOptimizerTest {
         mHandlerThread.quit();
         mThread.quit();
         mCountDown = null;
+        mDeviceConfig.tearDown();
     }
 
     private ProcessRecord makeProcessRecord(int pid, int uid, int packageUid, String processName,