diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipTaskOrganizerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipTaskOrganizerTest.java
index 800f9e4e5371ddd431884291b59aa037f2034530..75ee3f148c0cfad466aa60661394077c1f87812b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipTaskOrganizerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipTaskOrganizerTest.java
@@ -36,7 +36,6 @@ import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
 import android.graphics.Rect;
 import android.os.RemoteException;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.util.Rational;
@@ -45,6 +44,8 @@ import android.view.DisplayInfo;
 import android.view.SurfaceControl;
 import android.window.WindowContainerToken;
 
+import androidx.test.filters.SmallTest;
+
 import com.android.wm.shell.MockSurfaceControlHelper;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
index 4eb519334e12dd6a0f94d46c0bbb87c80f3b183c..a535274bd8bb1092a5fd74e10e7d9ba7dfd7dd7a 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
@@ -42,11 +42,12 @@ import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.RemoteException;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.util.Size;
 
+import androidx.test.filters.SmallTest;
+
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.WindowManagerShellWrapper;
 import com.android.wm.shell.common.DisplayController;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/tv/TvPipActionProviderTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/tv/TvPipActionProviderTest.java
index 45f6c8c7f69f1a1d454ab68f84b85500c4a98f1b..72db6e091307754d76cfc59eeb8e3bb361bf7671 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/tv/TvPipActionProviderTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/tv/TvPipActionProviderTest.java
@@ -23,20 +23,21 @@ import static com.android.wm.shell.pip.tv.TvPipAction.ACTION_EXPAND_COLLAPSE;
 import static com.android.wm.shell.pip.tv.TvPipAction.ACTION_FULLSCREEN;
 import static com.android.wm.shell.pip.tv.TvPipAction.ACTION_MOVE;
 
-import static java.util.Collections.EMPTY_LIST;
-
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import static java.util.Collections.EMPTY_LIST;
+
 import android.app.PendingIntent;
 import android.app.RemoteAction;
 import android.graphics.drawable.Icon;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
+import androidx.test.filters.SmallTest;
+
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.pip.PipMediaController;
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
index 050443914355c0016f16afc36e5935dc2f66085d..fbc0db9c28500fcde1869eae529287587d469934 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
@@ -32,7 +32,6 @@ import static org.mockito.Mockito.when;
 import android.app.ActivityManager;
 import android.graphics.Rect;
 import android.os.IBinder;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.SurfaceControl;
@@ -40,6 +39,8 @@ import android.window.TransitionInfo;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
+import androidx.test.filters.SmallTest;
+
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.transition.Transitions;