[ImsMedia] flaky test failure in ImsMediaJavaUnitTests

Flaky test failure due to different looper being used. Logic is modified
to use the same looper while performing UT.

Bug: 275030625
Test: atest -c ImsMediaJavaUnitTests --iterations 100
      Basic functionality is verified manually
Change-Id: Ie17c820c8454117b109796c909e9b387da8e7241

11 files changed, 165 insertions, 150 deletions

diff --git a/service/src/com/android/telephony/imsmedia/AudioSession.java b/service/src/com/android/telephony/imsmedia/AudioSession.java
index 603780c2..dae20c41 100644
--- a/service/src/com/android/telephony/imsmedia/AudioSession.java
+++ b/service/src/com/android/telephony/imsmedia/AudioSession.java
@@ -103,13 +103,14 @@ public final class AudioSession extends IImsAudioSession.Stub implements IMediaS
 
     @VisibleForTesting
     AudioSession(final int sessionId,
-            final @NonNull IImsAudioSessionCallback callback,
-            final @Nullable AudioService audioService,
-            final @Nullable AudioLocalSession localSession,
-            final @Nullable AudioOffloadService offloadService) {
+            @NonNull final IImsAudioSessionCallback callback,
+            @Nullable final AudioService audioService,
+            @Nullable final AudioLocalSession localSession,
+            @Nullable final AudioOffloadService offloadService,
+            Looper looper) {
         mSessionId = sessionId;
         mCallback = callback;
-        mHandler = new AudioSessionHandler(Looper.getMainLooper());
+        mHandler = new AudioSessionHandler(looper);
         mAudioService = audioService;
         mLocalSession = localSession;
         mAudioListener = new AudioListener(mHandler);
diff --git a/service/src/com/android/telephony/imsmedia/TextSession.java b/service/src/com/android/telephony/imsmedia/TextSession.java
index b0895d8e..dd069a03 100644
--- a/service/src/com/android/telephony/imsmedia/TextSession.java
+++ b/service/src/com/android/telephony/imsmedia/TextSession.java
@@ -73,10 +73,10 @@ public final class TextSession extends IImsTextSession.Stub implements IMediaSes
     TextSession(final int sessionId,
             final @NonNull IImsTextSessionCallback callback,
             final @Nullable TextService textService,
-            final @Nullable TextLocalSession localSession) {
+            final @Nullable TextLocalSession localSession, Looper looper) {
         mSessionId = sessionId;
         mCallback = callback;
-        mHandler = new TextSessionHandler(Looper.getMainLooper());
+        mHandler = new TextSessionHandler(looper);
         mTextService = textService;
         mLocalSession = localSession;
         mTextListener = new TextListener(mHandler);
diff --git a/service/src/com/android/telephony/imsmedia/VideoSession.java b/service/src/com/android/telephony/imsmedia/VideoSession.java
index a2351df9..5ade826e 100644
--- a/service/src/com/android/telephony/imsmedia/VideoSession.java
+++ b/service/src/com/android/telephony/imsmedia/VideoSession.java
@@ -85,10 +85,10 @@ public final class VideoSession extends IImsVideoSession.Stub implements IMediaS
     VideoSession(final int sessionId,
             final @NonNull IImsVideoSessionCallback callback,
             final @Nullable VideoService videoService,
-            final @Nullable VideoLocalSession localSession) {
+            final @Nullable VideoLocalSession localSession, Looper looper) {
         mSessionId = sessionId;
         mCallback = callback;
-        mHandler = new VideoSessionHandler(Looper.getMainLooper());
+        mHandler = new VideoSessionHandler(looper);
         mVideoService = videoService;
         mLocalSession = localSession;
         mVideoListener = new VideoListener(mHandler);
diff --git a/tests/unit/src/com/android/telephony/imsmedia/AudioListenerTest.java b/tests/unit/src/com/android/telephony/imsmedia/AudioListenerTest.java
index 353b6f7c..5dace993 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/AudioListenerTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/AudioListenerTest.java
@@ -22,6 +22,7 @@ import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.os.Looper;
 import android.os.Parcel;
 import android.os.RemoteException;
 import android.telephony.CallQuality;
@@ -30,21 +31,22 @@ import android.telephony.imsmedia.AudioConfig;
 import android.telephony.imsmedia.IImsAudioSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.MediaQualityStatus;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(JUnit4.class)
-public class AudioListenerTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class AudioListenerTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final char DTMF_DIGIT = '7';
     private static final int DTMF_DURATION = 120;
@@ -58,30 +60,23 @@ public class AudioListenerTest {
     @Mock
     private IImsAudioSessionCallback mMockIImsAudioSessionCallback;
     private AudioConfig mAudioConfig;
-    private TestableLooper mLooper;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         AudioSession audioSession = new AudioSession(SESSION_ID, mMockIImsAudioSessionCallback,
-                mAudioService, mMockAudioLocalSession, null);
+                mAudioService, mMockAudioLocalSession, null, Looper.myLooper());
         AudioSession.AudioSessionHandler handler = audioSession.getAudioSessionHandler();
         mAudioListener = new AudioListener(handler);
         mAudioListener.setMediaCallback(mMockCallback);
         mAudioConfig = AudioConfigTest.createAudioConfig();
-        try {
-            mLooper = new TestableLooper(handler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = AudioListenerTest.this;
+        super.setUp();
     }
 
     @After
     public void tearDown() throws Exception {
-        if (mLooper != null) {
-            mLooper.destroy();
-            mLooper = null;
-        }
+        super.tearDown();
     }
 
     private Parcel createParcel(int event, int result, AudioConfig config) {
@@ -238,10 +233,4 @@ public class AudioListenerTest {
         doNothing().when(mMockCallback).onSessionClosed(eq(SESSION_ID));
         verify(mMockCallback, times(1)).onSessionClosed(eq(SESSION_ID));
     }
-
-    private void processAllMessages() {
-        while (!mLooper.getLooper().getQueue().isIdle()) {
-            mLooper.processAllMessages();
-        }
-    }
 }
diff --git a/tests/unit/src/com/android/telephony/imsmedia/AudioOffloadTest.java b/tests/unit/src/com/android/telephony/imsmedia/AudioOffloadTest.java
index 5eb92175..23ec9b0d 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/AudioOffloadTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/AudioOffloadTest.java
@@ -30,6 +30,7 @@ import android.hardware.radio.ims.media.IImsMedia;
 import android.hardware.radio.ims.media.IImsMediaSession;
 import android.hardware.radio.ims.media.RtpConfig;
 import android.hardware.radio.ims.media.RtpError;
+import android.os.Looper;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
 import android.telephony.CallQuality;
@@ -39,6 +40,7 @@ import android.telephony.imsmedia.IImsAudioSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.MediaQualityStatus;
 import android.telephony.imsmedia.MediaQualityThreshold;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import com.android.telephony.imsmedia.AudioSession;
@@ -49,7 +51,6 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -60,8 +61,9 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
 
-@RunWith(JUnit4.class)
-public class AudioOffloadTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class AudioOffloadTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final int DTMF_DURATION = 120;
     private static final int NO_RESOURCES = ImsMediaSession.RESULT_NO_RESOURCES;
@@ -73,7 +75,6 @@ public class AudioOffloadTest {
     private AudioSession audioSession;
     private AudioOffloadListener offloadListener;
     private AudioSession.AudioSessionHandler handler;
-    private TestableLooper looper;
     @Mock
     private IImsAudioSessionCallback callback;
     @Mock
@@ -88,24 +89,19 @@ public class AudioOffloadTest {
         MockitoAnnotations.initMocks(this);
         offloadService = spy(AudioOffloadService.getInstance());
         doReturn(imsMedia).when(offloadService).getIImsMedia();
-        audioSession = new AudioSession(SESSION_ID, callback, null, null, offloadService);
+        audioSession = new AudioSession(SESSION_ID, callback, null, null, offloadService,
+                Looper.myLooper());
         handler = audioSession.getAudioSessionHandler();
         audioSession.setAudioOffload(true);
         offloadListener = audioSession.getOffloadListener();
         audioSession.onOpenSessionSuccess(imsMediaSession);
-        try {
-            looper = new TestableLooper(handler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = AudioOffloadTest.this;
+        super.setUp();
     }
 
     @After
     public void tearDown() throws Exception {
-        if (looper != null) {
-            looper.destroy();
-            looper = null;
-        }
+        super.tearDown();
     }
 
     @Test
@@ -397,10 +393,4 @@ public class AudioOffloadTest {
             fail("Failed to notify onCallQualityChanged: " + e);
         }
     }
-
-    private void processAllMessages() {
-        while (!looper.getLooper().getQueue().isIdle()) {
-            looper.processAllMessages();
-        }
-    }
 }
diff --git a/tests/unit/src/com/android/telephony/imsmedia/AudioSessionTest.java b/tests/unit/src/com/android/telephony/imsmedia/AudioSessionTest.java
index 2c9bf484..e51a46fb 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/AudioSessionTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/AudioSessionTest.java
@@ -21,6 +21,7 @@ import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.os.Looper;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
@@ -31,6 +32,7 @@ import android.telephony.imsmedia.IImsAudioSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.MediaQualityStatus;
 import android.telephony.imsmedia.MediaQualityThreshold;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import com.android.telephony.imsmedia.AudioService;
@@ -42,7 +44,6 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -50,8 +51,9 @@ import java.net.DatagramSocket;
 import java.net.SocketException;
 import java.util.ArrayList;
 
-@RunWith(JUnit4.class)
-public class AudioSessionTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class AudioSessionTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final int DTMF_DURATION = 140;
     private static final int UNUSED = -1;
@@ -69,28 +71,21 @@ public class AudioSessionTest {
     private AudioLocalSession audioLocalSession;
     @Mock
     private IImsAudioSessionCallback callback;
-    private TestableLooper looper;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         audioSession = new AudioSession(SESSION_ID, callback,
-                audioService, audioLocalSession, null);
+                audioService, audioLocalSession, null, Looper.myLooper());
         audioListener = audioSession.getAudioListener();
         handler = audioSession.getAudioSessionHandler();
-        try {
-            looper = new TestableLooper(handler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = AudioSessionTest.this;
+        super.setUp();
     }
 
     @After
     public void tearDown() throws Exception {
-        if (looper != null) {
-            looper.destroy();
-            looper = null;
-        }
+        super.tearDown();
     }
 
     private Parcel createParcel(int message, int result, AudioConfig config) {
@@ -366,10 +361,4 @@ public class AudioSessionTest {
             fail("Failed to notify onSessionClosed: " + e);
         }
     }
-
-    private void processAllMessages() {
-        while (!looper.getLooper().getQueue().isIdle()) {
-            looper.processAllMessages();
-        }
-    }
 }
diff --git a/tests/unit/src/com/android/telephony/imsmedia/ImsMediaTest.java b/tests/unit/src/com/android/telephony/imsmedia/ImsMediaTest.java
new file mode 100644
index 00000000..759e6f32
--- /dev/null
+++ b/tests/unit/src/com/android/telephony/imsmedia/ImsMediaTest.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.telephony.imsmedia;
+
+import static org.junit.Assert.fail;
+
+import android.testing.TestableLooper;
+
+import org.junit.After;
+import org.junit.Before;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class ImsMediaTest {
+    protected List<TestableLooper> mTestableLoopers = new ArrayList<>();
+    protected TestableLooper mTestableLooper;
+    protected Object mTestClass;
+
+    @Before
+    public void setUp() {
+        mTestableLooper = TestableLooper.get(mTestClass);
+        if (mTestableLooper != null) {
+            monitorTestableLooper(mTestableLooper);
+        }
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        if (!mTestableLoopers.isEmpty()) {
+            for (TestableLooper looper : mTestableLoopers) {
+                looper.getLooper().quit();
+            }
+        }
+        // Unmonitor TestableLooper for ImsMediaTest class
+        if (mTestableLooper != null) {
+            unmonitorTestableLooper(mTestableLooper);
+        }
+        // Destroy all newly created TestableLoopers so they can be reused
+        for (TestableLooper looper : mTestableLoopers) {
+            looper.destroy();
+        }
+        TestableLooper.remove(mTestClass);
+
+    }
+
+    private void monitorTestableLooper(TestableLooper looper) {
+        if (!mTestableLoopers.contains(looper)) {
+            mTestableLoopers.add(looper);
+        }
+    }
+
+    private void unmonitorTestableLooper(TestableLooper looper) {
+        if (mTestableLoopers.contains(looper)) {
+            mTestableLoopers.remove(looper);
+        }
+    }
+
+    private boolean areAllTestableLoopersIdle() {
+        for (TestableLooper looper : mTestableLoopers) {
+            if (!looper.getLooper().getQueue().isIdle()) return false;
+        }
+        return true;
+    }
+
+    public void processAllMessages() {
+        if (mTestableLoopers.isEmpty()) {
+            fail("mTestableLoopers is empty. Please make sure to add @RunWithLooper annotation");
+        }
+        while (!areAllTestableLoopersIdle()) {
+            for (TestableLooper looper : mTestableLoopers) looper.processAllMessages();
+        }
+    }
+}
diff --git a/tests/unit/src/com/android/telephony/imsmedia/TextListenerTest.java b/tests/unit/src/com/android/telephony/imsmedia/TextListenerTest.java
index ce2520c1..a7f016f5 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/TextListenerTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/TextListenerTest.java
@@ -22,23 +22,25 @@ import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.os.Looper;
 import android.os.Parcel;
 import android.os.RemoteException;
 import android.telephony.imsmedia.IImsTextSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.TextConfig;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-@RunWith(JUnit4.class)
-public class TextListenerTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class TextListenerTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final String TEXT_STREAM = "Hello";
     private static final long NATIVE_OBJECT = 1234L;
@@ -52,31 +54,24 @@ public class TextListenerTest {
     @Mock
     private IImsTextSessionCallback mMockIImsTextSessionCallback;
     private TextConfig mTextConfig;
-    private TestableLooper mLooper;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         TextSession textSession = new TextSession(SESSION_ID, mMockIImsTextSessionCallback,
-                mTextService, mMockTextLocalSession);
+                mTextService, mMockTextLocalSession, Looper.myLooper());
         TextSession.TextSessionHandler handler = textSession.getTextSessionHandler();
         mTextListener = new TextListener(handler);
         mTextListener.setMediaCallback(mMockCallback);
         mTextListener.setNativeObject(NATIVE_OBJECT);
         mTextConfig = TextConfigTest.createTextConfig();
-        try {
-            mLooper = new TestableLooper(handler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = TextListenerTest.this;
+        super.setUp();
     }
 
     @After
-    public void tearDown() {
-        if (mLooper != null) {
-            mLooper.destroy();
-            mLooper = null;
-        }
+    public void tearDown() throws Exception {
+        super.tearDown();
     }
 
     private Parcel createParcel(int event, int result, TextConfig config) {
@@ -166,10 +161,4 @@ public class TextListenerTest {
         parcel.recycle();
         verify(mMockCallback, times(1)).onSessionClosed(eq(SESSION_ID));
     }
-
-    private void processAllMessages() {
-        while (!mLooper.getLooper().getQueue().isIdle()) {
-            mLooper.processAllMessages();
-        }
-    }
 }
diff --git a/tests/unit/src/com/android/telephony/imsmedia/TextSessionTest.java b/tests/unit/src/com/android/telephony/imsmedia/TextSessionTest.java
index 7319711e..0c150026 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/TextSessionTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/TextSessionTest.java
@@ -21,6 +21,7 @@ import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.os.Looper;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
@@ -28,6 +29,7 @@ import android.telephony.imsmedia.IImsTextSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.MediaQualityThreshold;
 import android.telephony.imsmedia.TextConfig;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import com.android.telephony.imsmedia.Utils.OpenSessionParams;
@@ -36,15 +38,15 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.net.DatagramSocket;
 import java.net.SocketException;
 
-@RunWith(JUnit4.class)
-public class TextSessionTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class TextSessionTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final int SUCCESS = ImsMediaSession.RESULT_SUCCESS;
     private static final int NO_RESOURCES = ImsMediaSession.RESULT_NO_RESOURCES;
@@ -60,27 +62,21 @@ public class TextSessionTest {
     private TextLocalSession mTextLocalSession;
     @Mock
     private IImsTextSessionCallback mCallback;
-    private TestableLooper mLooper;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mTextSession = new TextSession(SESSION_ID, mCallback, mTextService, mTextLocalSession);
+        mTextSession = new TextSession(SESSION_ID, mCallback, mTextService, mTextLocalSession,
+                Looper.myLooper());
         mTextListener = mTextSession.getTextListener();
         mHandler = mTextSession.getTextSessionHandler();
-        try {
-            mLooper = new TestableLooper(mHandler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = TextSessionTest.this;
+        super.setUp();
     }
 
     @After
     public void tearDown() throws Exception {
-        if (mLooper != null) {
-            mLooper.destroy();
-            mLooper = null;
-        }
+        super.tearDown();
     }
 
     private Parcel createParcel(int message, int result, TextConfig config) {
@@ -241,10 +237,4 @@ public class TextSessionTest {
             fail("Failed to notify onSessionClosed: " + e);
         }
     }
-
-    private void processAllMessages() {
-        while (!mLooper.getLooper().getQueue().isIdle()) {
-            mLooper.processAllMessages();
-        }
-    }
 }
diff --git a/tests/unit/src/com/android/telephony/imsmedia/VideoListenerTest.java b/tests/unit/src/com/android/telephony/imsmedia/VideoListenerTest.java
index ad90f957..a9948d22 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/VideoListenerTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/VideoListenerTest.java
@@ -22,23 +22,25 @@ import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.os.Looper;
 import android.os.Parcel;
 import android.os.RemoteException;
 import android.telephony.imsmedia.IImsVideoSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.VideoConfig;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-@RunWith(JUnit4.class)
-public class VideoListenerTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class VideoListenerTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final long VIDEO_DATA = 1024;
     private static final int RESOLUTION_WIDTH = 640;
@@ -54,31 +56,24 @@ public class VideoListenerTest {
     @Mock
     private IImsVideoSessionCallback mMockIImsVideoSessionCallback;
     private VideoConfig mVideoConfig;
-    private TestableLooper mLooper;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         VideoSession videoSession = new VideoSession(SESSION_ID, mMockIImsVideoSessionCallback,
-                mVideoService, mMockVideoLocalSession);
+                mVideoService, mMockVideoLocalSession, Looper.myLooper());
         VideoSession.VideoSessionHandler handler = videoSession.getVideoSessionHandler();
         mVideoListener = new VideoListener(handler);
         mVideoListener.setMediaCallback(mMockCallback);
         mVideoListener.setNativeObject(NATIVE_OBJECT);
         mVideoConfig = VideoConfigTest.createVideoConfig();
-        try {
-            mLooper = new TestableLooper(handler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = VideoListenerTest.this;
+        super.setUp();
     }
 
     @After
-    public void tearDown() {
-        if (mLooper != null) {
-            mLooper.destroy();
-            mLooper = null;
-        }
+    public void tearDown() throws Exception {
+        super.tearDown();
     }
 
     private Parcel createParcel(int event, int result) {
@@ -201,10 +196,4 @@ public class VideoListenerTest {
         parcel.recycle();
         verify(mMockCallback, times(1)).onSessionClosed(eq(SESSION_ID));
     }
-
-    private void processAllMessages() {
-        while (!mLooper.getLooper().getQueue().isIdle()) {
-            mLooper.processAllMessages();
-        }
-    }
 }
diff --git a/tests/unit/src/com/android/telephony/imsmedia/VideoSessionTest.java b/tests/unit/src/com/android/telephony/imsmedia/VideoSessionTest.java
index 3f781886..600f9778 100644
--- a/tests/unit/src/com/android/telephony/imsmedia/VideoSessionTest.java
+++ b/tests/unit/src/com/android/telephony/imsmedia/VideoSessionTest.java
@@ -23,6 +23,7 @@ import static org.mockito.Mockito.verify;
 
 import android.graphics.ImageFormat;
 import android.media.ImageReader;
+import android.os.Looper;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
@@ -31,6 +32,7 @@ import android.telephony.imsmedia.IImsVideoSessionCallback;
 import android.telephony.imsmedia.ImsMediaSession;
 import android.telephony.imsmedia.MediaQualityThreshold;
 import android.telephony.imsmedia.VideoConfig;
+import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.Surface;
 
@@ -40,7 +42,6 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -48,8 +49,9 @@ import java.net.DatagramSocket;
 import java.net.SocketException;
 import java.util.ArrayList;
 
-@RunWith(JUnit4.class)
-public class VideoSessionTest {
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class VideoSessionTest extends ImsMediaTest {
     private static final int SESSION_ID = 1;
     private static final int UNUSED = -1;
     private static final int SUCCESS = ImsMediaSession.RESULT_SUCCESS;
@@ -79,22 +81,16 @@ public class VideoSessionTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mVideoSession = new VideoSession(SESSION_ID, mCallback,
-                mVideoService, mVideoLocalSession);
+                mVideoService, mVideoLocalSession, Looper.myLooper());
         mVideoListener = mVideoSession.getVideoListener();
         mHandler = mVideoSession.getVideoSessionHandler();
-        try {
-            mLooper = new TestableLooper(mHandler.getLooper());
-        } catch (Exception e) {
-            throw new AssertionError("Unable to create TestableLooper", e);
-        }
+        mTestClass = VideoSessionTest.this;
+        super.setUp();
     }
 
     @After
     public void tearDown() throws Exception {
-        if (mLooper != null) {
-            mLooper.destroy();
-            mLooper = null;
-        }
+        super.tearDown();
     }
 
     private Parcel createParcel(int message, int result, VideoConfig config) {
@@ -330,10 +326,4 @@ public class VideoSessionTest {
             fail("Failed to notify onSessionClosed: " + e);
         }
     }
-
-    private void processAllMessages() {
-        while (!mLooper.getLooper().getQueue().isIdle()) {
-            mLooper.processAllMessages();
-        }
-    }
 }