diff --git a/tests/functional/com/android/documentsui/archives/WriteableArchiveTest.java b/tests/functional/com/android/documentsui/archives/WriteableArchiveTest.java
index 7092d12..da470da 100644
--- a/tests/functional/com/android/documentsui/archives/WriteableArchiveTest.java
+++ b/tests/functional/com/android/documentsui/archives/WriteableArchiveTest.java
@@ -21,9 +21,9 @@
 import android.os.ParcelFileDescriptor;
 import android.provider.DocumentsContract.Document;
 import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.MediumTest;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.MediumTest;
 
 import java.io.File;
 import java.io.IOException;
diff --git a/tests/functional/com/android/documentsui/inspector/DocumentLoaderTest.java b/tests/functional/com/android/documentsui/inspector/DocumentLoaderTest.java
index 8a2b21f..ac417c9 100644
--- a/tests/functional/com/android/documentsui/inspector/DocumentLoaderTest.java
+++ b/tests/functional/com/android/documentsui/inspector/DocumentLoaderTest.java
@@ -25,8 +25,8 @@
 import android.os.Bundle;
 import android.os.Looper;
 import android.provider.DocumentsContract;
-import android.test.suitebuilder.annotation.MediumTest;
 
+import androidx.test.filters.MediumTest;
 import androidx.test.rule.provider.ProviderTestRule;
 
 import com.android.documentsui.InspectorProvider;
diff --git a/tests/functional/com/android/documentsui/services/AbstractCopyJobTest.java b/tests/functional/com/android/documentsui/services/AbstractCopyJobTest.java
index 7b2cb82..525397b 100644
--- a/tests/functional/com/android/documentsui/services/AbstractCopyJobTest.java
+++ b/tests/functional/com/android/documentsui/services/AbstractCopyJobTest.java
@@ -23,9 +23,10 @@
 import android.app.Notification;
 import android.net.Uri;
 import android.provider.DocumentsContract;
-import android.test.suitebuilder.annotation.MediumTest;
 import android.text.format.DateUtils;
 
+import androidx.test.filters.MediumTest;
+
 import com.android.documentsui.R;
 import com.android.documentsui.base.DocumentInfo;
 import com.android.documentsui.services.FileOperationService.OpType;
diff --git a/tests/unit/com/android/documentsui/FocusManagerTest.java b/tests/unit/com/android/documentsui/FocusManagerTest.java
index 67e7000..36637f3 100644
--- a/tests/unit/com/android/documentsui/FocusManagerTest.java
+++ b/tests/unit/com/android/documentsui/FocusManagerTest.java
@@ -17,10 +17,10 @@
 package com.android.documentsui;
 
 import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
 
 import androidx.recyclerview.selection.SelectionTracker;
 import androidx.recyclerview.widget.RecyclerView;
+import androidx.test.filters.SmallTest;
 
 import com.android.documentsui.base.UserId;
 import com.android.documentsui.dirlist.TestData;
diff --git a/tests/unit/com/android/documentsui/base/DocumentInfoTest.java b/tests/unit/com/android/documentsui/base/DocumentInfoTest.java
index eba30d2..4e73333 100644
--- a/tests/unit/com/android/documentsui/base/DocumentInfoTest.java
+++ b/tests/unit/com/android/documentsui/base/DocumentInfoTest.java
@@ -28,8 +28,8 @@
 import android.net.Uri;
 import android.provider.DocumentsContract;
 import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
 
+import androidx.test.filters.SmallTest;
 import androidx.test.rule.provider.ProviderTestRule;
 
 import com.android.documentsui.InspectorProvider;
diff --git a/tests/unit/com/android/documentsui/dirlist/AccessibilityTest.java b/tests/unit/com/android/documentsui/dirlist/AccessibilityTest.java
index 4334608..09b1933 100644
--- a/tests/unit/com/android/documentsui/dirlist/AccessibilityTest.java
+++ b/tests/unit/com/android/documentsui/dirlist/AccessibilityTest.java
@@ -18,12 +18,12 @@
 
 import android.database.Cursor;
 import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.view.View;
 import android.widget.Space;
 
 import androidx.core.view.accessibility.AccessibilityNodeInfoCompat;
 import androidx.recyclerview.widget.RecyclerView;
+import androidx.test.filters.SmallTest;
 
 import com.android.documentsui.TestConfigStore;
 import com.android.documentsui.testing.TestRecyclerView;
diff --git a/tests/unit/com/android/documentsui/inspector/GpsCoordinatesTextClassifierTest.java b/tests/unit/com/android/documentsui/inspector/GpsCoordinatesTextClassifierTest.java
index e25d227..f4e0422 100644
--- a/tests/unit/com/android/documentsui/inspector/GpsCoordinatesTextClassifierTest.java
+++ b/tests/unit/com/android/documentsui/inspector/GpsCoordinatesTextClassifierTest.java
@@ -20,12 +20,12 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.os.LocaleList;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.view.textclassifier.TextClassification;
 import android.view.textclassifier.TextClassificationManager;
 import android.view.textclassifier.TextClassifier;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.documentsui.testing.TestPackageManager;
diff --git a/tests/unit/com/android/documentsui/inspector/HeaderTextSelectorTest.java b/tests/unit/com/android/documentsui/inspector/HeaderTextSelectorTest.java
index 30178ce..28bc667 100644
--- a/tests/unit/com/android/documentsui/inspector/HeaderTextSelectorTest.java
+++ b/tests/unit/com/android/documentsui/inspector/HeaderTextSelectorTest.java
@@ -19,12 +19,12 @@
 import static junit.framework.Assert.assertTrue;
 
 import android.content.Context;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.text.Spannable;
 import android.text.SpannableString;
 import android.widget.TextView;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.documentsui.inspector.HeaderTextSelector.Selector;
diff --git a/tests/unit/com/android/documentsui/inspector/InspectorControllerTest.java b/tests/unit/com/android/documentsui/inspector/InspectorControllerTest.java
index 9a7bad5..f3d5b59 100644
--- a/tests/unit/com/android/documentsui/inspector/InspectorControllerTest.java
+++ b/tests/unit/com/android/documentsui/inspector/InspectorControllerTest.java
@@ -29,10 +29,10 @@
 import android.os.UserHandle;
 import android.provider.DocumentsContract;
 import android.provider.DocumentsContract.Document;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.view.View.OnClickListener;
 
 import androidx.annotation.Nullable;
+import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.documentsui.InspectorProvider;
diff --git a/tests/unit/com/android/documentsui/inspector/MediaViewTest.java b/tests/unit/com/android/documentsui/inspector/MediaViewTest.java
index 80c3c47..db87e42 100644
--- a/tests/unit/com/android/documentsui/inspector/MediaViewTest.java
+++ b/tests/unit/com/android/documentsui/inspector/MediaViewTest.java
@@ -19,8 +19,8 @@
 import android.media.MediaMetadata;
 import android.os.Bundle;
 import android.provider.DocumentsContract;
-import android.test.suitebuilder.annotation.SmallTest;
 
+import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.documentsui.R;
diff --git a/tests/unit/com/android/documentsui/roots/ProvidersAccessTest.java b/tests/unit/com/android/documentsui/roots/ProvidersAccessTest.java
index cab6641..255c1cb 100644
--- a/tests/unit/com/android/documentsui/roots/ProvidersAccessTest.java
+++ b/tests/unit/com/android/documentsui/roots/ProvidersAccessTest.java
@@ -21,9 +21,9 @@
 
 import android.provider.DocumentsContract;
 import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
 
 import androidx.annotation.Nullable;
+import androidx.test.filters.SmallTest;
 
 import com.android.documentsui.base.Providers;
 import com.android.documentsui.base.RootInfo;