diff --git a/tests/SurfaceComposition/Android.bp b/tests/SurfaceComposition/Android.bp
index f5aba8f5a2f2e94200637c6fb154a27cf85ba5a9..a02662fb912ec983eca8b5a3f865b35009c6d0bd 100644
--- a/tests/SurfaceComposition/Android.bp
+++ b/tests/SurfaceComposition/Android.bp
@@ -32,7 +32,10 @@ android_test {
         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 261ea2ec866d568d837335ee44b5a91bc0d553c3..f82585c02c08dfaff89b01e5664dd5ecaec433f5 100644
--- a/tests/SurfaceComposition/src/android/surfacecomposition/SurfaceCompositionTest.java
+++ b/tests/SurfaceComposition/src/android/surfacecomposition/SurfaceCompositionTest.java
@@ -22,9 +22,10 @@ import android.os.Bundle;
 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";