diff --git a/test-runner/tests/Android.bp b/test-runner/tests/Android.bp
index aad2bee..0c0c080 100644
--- a/test-runner/tests/Android.bp
+++ b/test-runner/tests/Android.bp
@@ -33,7 +33,10 @@
         "android.test.base",
         "android.test.mock",
     ],
-    static_libs: ["junit"],
+    static_libs: [
+        "junit",
+        "androidx.test.rules",
+    ],
 
     // Include all test java files.
     srcs: ["src/**/*.java"],
diff --git a/test-runner/tests/src/android/test/AndroidTestRunnerTest.java b/test-runner/tests/src/android/test/AndroidTestRunnerTest.java
index bd6c04bc..9b7ea59 100644
--- a/test-runner/tests/src/android/test/AndroidTestRunnerTest.java
+++ b/test-runner/tests/src/android/test/AndroidTestRunnerTest.java
@@ -17,7 +17,8 @@
 package android.test;
 
 import android.test.mock.MockContext;
-import android.test.suitebuilder.annotation.SmallTest;
+
+import androidx.test.filters.SmallTest;
 
 import junit.framework.AssertionFailedError;
 import junit.framework.Test;
diff --git a/test-runner/tests/src/android/test/InstrumentationTestRunnerTest.java b/test-runner/tests/src/android/test/InstrumentationTestRunnerTest.java
index d98b217..0929e8b 100644
--- a/test-runner/tests/src/android/test/InstrumentationTestRunnerTest.java
+++ b/test-runner/tests/src/android/test/InstrumentationTestRunnerTest.java
@@ -22,7 +22,8 @@
 import android.test.mock.MockContext;
 import android.test.suitebuilder.ListTestCaseNames;
 import android.test.suitebuilder.ListTestCaseNames.TestDescriptor;
-import android.test.suitebuilder.annotation.SmallTest;
+
+import androidx.test.filters.SmallTest;
 
 import junit.framework.Test;
 import junit.framework.TestCase;
diff --git a/test-runner/tests/src/android/test/suitebuilder/examples/suppress/PartiallySuppressedTest.java b/test-runner/tests/src/android/test/suitebuilder/examples/suppress/PartiallySuppressedTest.java
index 3ca0f70..f48443d 100644
--- a/test-runner/tests/src/android/test/suitebuilder/examples/suppress/PartiallySuppressedTest.java
+++ b/test-runner/tests/src/android/test/suitebuilder/examples/suppress/PartiallySuppressedTest.java
@@ -16,7 +16,7 @@
 
 package android.test.suitebuilder.examples.suppress;
 
-import android.test.suitebuilder.annotation.Suppress;
+import androidx.test.filters.Suppress;
 
 import junit.framework.TestCase;
 
diff --git a/test-runner/tests/src/android/test/suitebuilder/examples/suppress/SuppressedTest.java b/test-runner/tests/src/android/test/suitebuilder/examples/suppress/SuppressedTest.java
index c4e0e07..c85238c 100644
--- a/test-runner/tests/src/android/test/suitebuilder/examples/suppress/SuppressedTest.java
+++ b/test-runner/tests/src/android/test/suitebuilder/examples/suppress/SuppressedTest.java
@@ -16,7 +16,7 @@
 
 package android.test.suitebuilder.examples.suppress;
 
-import android.test.suitebuilder.annotation.Suppress;
+import androidx.test.filters.Suppress;
 
 import junit.framework.TestCase;