diff --git a/tests/SurfaceComposition/Android.bp b/tests/SurfaceComposition/Android.bp
index f5aba8f5..a02662f 100644
--- a/tests/SurfaceComposition/Android.bp
+++ b/tests/SurfaceComposition/Android.bp
@@ -32,7 +32,10 @@
         enabled: false,
     },
     srcs: ["src/**/*.java"],
-    static_libs: ["junit"],
+    static_libs: [
+        "junit",
+        "androidx.test.rules",
+    ],
     libs: [
         "android.test.runner.stubs",
         "android.test.base.stubs",
diff --git a/tests/SurfaceComposition/src/android/surfacecomposition/SurfaceCompositionTest.java b/tests/SurfaceComposition/src/android/surfacecomposition/SurfaceCompositionTest.java
index 261ea2e..f82585c 100644
--- a/tests/SurfaceComposition/src/android/surfacecomposition/SurfaceCompositionTest.java
+++ b/tests/SurfaceComposition/src/android/surfacecomposition/SurfaceCompositionTest.java
@@ -22,9 +22,10 @@
 import android.surfacecomposition.SurfaceCompositionMeasuringActivity.AllocationScore;
 import android.surfacecomposition.SurfaceCompositionMeasuringActivity.CompositorScore;
 import android.test.ActivityInstrumentationTestCase2;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.util.Log;
 
+import androidx.test.filters.SmallTest;
+
 public class SurfaceCompositionTest extends
         ActivityInstrumentationTestCase2<SurfaceCompositionMeasuringActivity> {
     private final static String TAG = "SurfaceCompositionTest";