Merge "Create VibratorInfo#areVibrationFeaturesSupported" into main
diff --git a/core/java/android/os/VibrationEffect.java b/core/java/android/os/VibrationEffect.java
index b24b45d..08b32bf 100644
--- a/core/java/android/os/VibrationEffect.java
+++ b/core/java/android/os/VibrationEffect.java
@@ -525,14 +525,14 @@
     public abstract long getDuration();
 
     /**
-     * Checks if a given {@link Vibrator} can play this effect as intended.
+     * Checks if a vibrator with a given {@link VibratorInfo} can play this effect as intended.
      *
-     * <p>See @link Vibrator#areVibrationFeaturesSupported(VibrationEffect)} for more information
-     * about what counts as supported by a vibrator, and what counts as not.
+     * <p>See {@link VibratorInfo#areVibrationFeaturesSupported(VibrationEffect)} for more
+     * information about what counts as supported by a vibrator, and what counts as not.
      *
      * @hide
      */
-    public abstract boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator);
+    public abstract boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo);
 
     /**
      * Returns true if this effect could represent a touch haptic feedback.
@@ -813,9 +813,9 @@
 
         /** @hide */
         @Override
-        public boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator) {
+        public boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo) {
             for (VibrationEffectSegment segment : mSegments) {
-                if (!segment.areVibrationFeaturesSupported(vibrator)) {
+                if (!segment.areVibrationFeaturesSupported(vibratorInfo)) {
                     return false;
                 }
             }
diff --git a/core/java/android/os/Vibrator.java b/core/java/android/os/Vibrator.java
index 4e852e3..79e0ca8 100644
--- a/core/java/android/os/Vibrator.java
+++ b/core/java/android/os/Vibrator.java
@@ -216,9 +216,7 @@
      */
     @TestApi
     public boolean hasFrequencyControl() {
-        // We currently can only control frequency of the vibration using the compose PWLE method.
-        return getInfo().hasCapability(
-                IVibrator.CAP_FREQUENCY_CONTROL | IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
+        return getInfo().hasFrequencyControl();
     }
 
     /**
@@ -240,7 +238,7 @@
      * @hide
      */
     public boolean areVibrationFeaturesSupported(@NonNull VibrationEffect effect) {
-        return effect.areVibrationFeaturesSupported(this);
+        return getInfo().areVibrationFeaturesSupported(effect);
     }
 
     /**
diff --git a/core/java/android/os/VibratorInfo.java b/core/java/android/os/VibratorInfo.java
index 02e6856..0b7d7c3 100644
--- a/core/java/android/os/VibratorInfo.java
+++ b/core/java/android/os/VibratorInfo.java
@@ -242,6 +242,17 @@
     }
 
     /**
+     * Check whether the vibrator has frequency control.
+     *
+     * @return True if the hardware can control the frequency of the vibrations, otherwise false.
+     */
+    public boolean hasFrequencyControl() {
+        // We currently can only control frequency of the vibration using the compose PWLE method.
+        return hasCapability(
+                IVibrator.CAP_FREQUENCY_CONTROL | IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
+    }
+
+    /**
      * Returns a default value to be applied to composed PWLE effects for braking.
      *
      * @return a supported braking value, one of android.hardware.vibrator.Braking.*
@@ -323,6 +334,23 @@
     }
 
     /**
+     * Query whether or not the vibrator supports all components of a given {@link VibrationEffect}
+     * (i.e. the vibrator can play the given effect as intended).
+     *
+     * <p>See {@link Vibrator#areVibrationFeaturesSupported(VibrationEffect)} for more
+     * information on how the vibrator support is determined.
+     *
+     * @param effect the {@link VibrationEffect} to check if it is supported
+     * @return {@code true} if the vibrator can play the given {@code effect} as intended,
+     *         {@code false} otherwise.
+     *
+     * @hide
+     */
+    public boolean areVibrationFeaturesSupported(@NonNull VibrationEffect effect) {
+        return effect.areVibrationFeaturesSupported(this);
+    }
+
+    /**
      * Query the estimated duration of given primitive.
      *
      * @param primitiveId Which primitives to query for.
diff --git a/core/java/android/os/vibrator/PrebakedSegment.java b/core/java/android/os/vibrator/PrebakedSegment.java
index 42b6c2da..a035092 100644
--- a/core/java/android/os/vibrator/PrebakedSegment.java
+++ b/core/java/android/os/vibrator/PrebakedSegment.java
@@ -23,6 +23,7 @@
 import android.os.Parcelable;
 import android.os.VibrationEffect;
 import android.os.Vibrator;
+import android.os.VibratorInfo;
 
 import java.util.Objects;
 
@@ -77,8 +78,8 @@
 
     /** @hide */
     @Override
-    public boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator) {
-        if (vibrator.areAllEffectsSupported(mEffectId) == Vibrator.VIBRATION_EFFECT_SUPPORT_YES) {
+    public boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo) {
+        if (vibratorInfo.isEffectSupported(mEffectId) == Vibrator.VIBRATION_EFFECT_SUPPORT_YES) {
             return true;
         }
         if (!mFallback) {
diff --git a/core/java/android/os/vibrator/PrimitiveSegment.java b/core/java/android/os/vibrator/PrimitiveSegment.java
index c52a09c..95d97bf 100644
--- a/core/java/android/os/vibrator/PrimitiveSegment.java
+++ b/core/java/android/os/vibrator/PrimitiveSegment.java
@@ -22,7 +22,7 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
+import android.os.VibratorInfo;
 
 import com.android.internal.util.Preconditions;
 
@@ -77,8 +77,8 @@
 
     /** @hide */
     @Override
-    public boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator) {
-        return vibrator.areAllPrimitivesSupported(mPrimitiveId);
+    public boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo) {
+        return vibratorInfo.isPrimitiveSupported(mPrimitiveId);
     }
 
     /** @hide */
diff --git a/core/java/android/os/vibrator/RampSegment.java b/core/java/android/os/vibrator/RampSegment.java
index e997bcd..5f9d102 100644
--- a/core/java/android/os/vibrator/RampSegment.java
+++ b/core/java/android/os/vibrator/RampSegment.java
@@ -20,7 +20,7 @@
 import android.annotation.TestApi;
 import android.os.Parcel;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
+import android.os.VibratorInfo;
 
 import com.android.internal.util.Preconditions;
 
@@ -96,7 +96,7 @@
 
     /** @hide */
     @Override
-    public boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator) {
+    public boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo) {
         boolean areFeaturesSupported = true;
         // If the start/end frequencies are not the same, require frequency control since we need to
         // ramp up/down the frequency.
@@ -104,7 +104,7 @@
                 // If there is no frequency ramping, make sure that the one frequency used does not
                 // require frequency control.
                 || frequencyRequiresFrequencyControl(mStartFrequencyHz)) {
-            areFeaturesSupported &= vibrator.hasFrequencyControl();
+            areFeaturesSupported &= vibratorInfo.hasFrequencyControl();
         }
         // If the start/end amplitudes are not the same, require amplitude control since we need to
         // ramp up/down the amplitude.
@@ -112,7 +112,7 @@
                 // If there is no amplitude ramping, make sure that the amplitude used does not
                 // require amplitude control.
                 || amplitudeRequiresAmplitudeControl(mStartAmplitude)) {
-            areFeaturesSupported &= vibrator.hasAmplitudeControl();
+            areFeaturesSupported &= vibratorInfo.hasAmplitudeControl();
         }
         return areFeaturesSupported;
     }
diff --git a/core/java/android/os/vibrator/StepSegment.java b/core/java/android/os/vibrator/StepSegment.java
index a585aa8..9576a5b 100644
--- a/core/java/android/os/vibrator/StepSegment.java
+++ b/core/java/android/os/vibrator/StepSegment.java
@@ -21,7 +21,7 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
+import android.os.VibratorInfo;
 
 import com.android.internal.util.Preconditions;
 
@@ -82,13 +82,13 @@
 
     /** @hide */
     @Override
-    public boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator) {
+    public boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo) {
         boolean areFeaturesSupported = true;
         if (frequencyRequiresFrequencyControl(mFrequencyHz)) {
-            areFeaturesSupported &= vibrator.hasFrequencyControl();
+            areFeaturesSupported &= vibratorInfo.hasFrequencyControl();
         }
         if (amplitudeRequiresAmplitudeControl(mAmplitude)) {
-            areFeaturesSupported &= vibrator.hasAmplitudeControl();
+            areFeaturesSupported &= vibratorInfo.hasAmplitudeControl();
         }
         return areFeaturesSupported;
     }
diff --git a/core/java/android/os/vibrator/VibrationEffectSegment.java b/core/java/android/os/vibrator/VibrationEffectSegment.java
index 3b286a7..17ac36f 100644
--- a/core/java/android/os/vibrator/VibrationEffectSegment.java
+++ b/core/java/android/os/vibrator/VibrationEffectSegment.java
@@ -21,7 +21,7 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
+import android.os.VibratorInfo;
 
 /**
  * Representation of a single segment of a {@link VibrationEffect}.
@@ -65,7 +65,7 @@
      *
      * @hide
      */
-    public abstract boolean areVibrationFeaturesSupported(@NonNull Vibrator vibrator);
+    public abstract boolean areVibrationFeaturesSupported(@NonNull VibratorInfo vibratorInfo);
 
     /**
      * Returns true if this segment could be a haptic feedback effect candidate.
diff --git a/core/tests/vibrator/src/android/os/VibrationEffectTest.java b/core/tests/vibrator/src/android/os/VibrationEffectTest.java
index 8be489e..e875875 100644
--- a/core/tests/vibrator/src/android/os/VibrationEffectTest.java
+++ b/core/tests/vibrator/src/android/os/VibrationEffectTest.java
@@ -41,8 +41,6 @@
 import android.os.vibrator.PrimitiveSegment;
 import android.os.vibrator.StepSegment;
 
-import androidx.test.InstrumentationRegistry;
-
 import com.android.internal.R;
 
 import org.jetbrains.annotations.NotNull;
@@ -838,31 +836,29 @@
 
     @Test
     public void testAreVibrationFeaturesSupported_allSegmentsSupported() {
-        Vibrator vibrator =
-                createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                        .setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL)
-                        .build());
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL)
+                .build();
 
         assertTrue(VibrationEffect.createWaveform(
                         /* timings= */ new long[] {1, 2, 3}, /* repeatIndex= */ -1)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(VibrationEffect.createWaveform(
                         /* timings= */ new long[] {1, 2, 3},
                         /* amplitudes= */ new int[] {10, 20, 40},
                         /* repeatIndex= */ 2)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(
                 VibrationEffect.startComposition()
                         .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_CLICK))
                         .repeatEffectIndefinitely(TEST_ONE_SHOT)
                         .compose()
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testAreVibrationFeaturesSupported_withUnsupportedSegments() {
-        Vibrator vibrator =
-                createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1).build());
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1).build();
 
         assertFalse(
                 VibrationEffect.startComposition()
@@ -872,7 +868,7 @@
                                 /* amplitudes= */ new int[] {10, 20, 40},
                                 /* repeatIndex= */ -1))
                         .compose()
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
@@ -996,13 +992,4 @@
 
         return context;
     }
-
-    private Vibrator createVibratorWithCustomInfo(VibratorInfo info) {
-        return new SystemVibrator(InstrumentationRegistry.getContext()) {
-            @Override
-            public VibratorInfo getInfo() {
-                return info;
-            }
-        };
-    }
 }
diff --git a/core/tests/vibrator/src/android/os/VibratorInfoTest.java b/core/tests/vibrator/src/android/os/VibratorInfoTest.java
index ff917aa..808c4ec 100644
--- a/core/tests/vibrator/src/android/os/VibratorInfoTest.java
+++ b/core/tests/vibrator/src/android/os/VibratorInfoTest.java
@@ -57,6 +57,17 @@
     }
 
     @Test
+    public void testHasFrequencyControl() {
+        VibratorInfo noCapabilities = new VibratorInfo.Builder(TEST_VIBRATOR_ID).build();
+        assertFalse(noCapabilities.hasFrequencyControl());
+        VibratorInfo composeAndFrequencyControl = new VibratorInfo.Builder(TEST_VIBRATOR_ID)
+                .setCapabilities(
+                        IVibrator.CAP_FREQUENCY_CONTROL | IVibrator.CAP_COMPOSE_PWLE_EFFECTS)
+                .build();
+        assertTrue(composeAndFrequencyControl.hasFrequencyControl());
+    }
+
+    @Test
     public void testHasCapabilities() {
         VibratorInfo info = new VibratorInfo.Builder(TEST_VIBRATOR_ID)
                 .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
@@ -335,4 +346,186 @@
         assertEquals(original, restored);
         assertEquals(original.hashCode(), restored.hashCode());
     }
+
+    @Test
+    public void areVibrationFeaturesSupported_noAmplitudeControl_fractionalAmplitudeUnsupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1).build();
+
+        // Have at least one fractional amplitude (amplitude not min (0) or max (255) or DEFAULT).
+        assertFalse(info.areVibrationFeaturesSupported(waveformWithAmplitudes(10, 30)));
+        assertFalse(info.areVibrationFeaturesSupported(waveformWithAmplitudes(10, 255)));
+        assertFalse(info.areVibrationFeaturesSupported(
+                VibrationEffect.createOneShot(20, /* amplitude= */ 40)));
+    }
+
+    @Test
+    public void areVibrationFeaturesSupported_noAmplitudeControl_nonFractionalAmplitudeSupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1).build();
+
+        // All amplitudes are min, max, or default. Requires no amplitude control.
+        assertTrue(info.areVibrationFeaturesSupported(
+                waveformWithAmplitudes(255, 0, VibrationEffect.DEFAULT_AMPLITUDE, 255)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createWaveform(
+                        /* timings= */ new long[] {1, 2, 3}, /* repeatIndex= */ -1)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createOneShot(20, VibrationEffect.DEFAULT_AMPLITUDE)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createOneShot(20, /* amplitude= */ 255)));
+    }
+
+    @Test
+    public void areVibrationFeaturesSupported_withAmplitudeControl_allWaveformsSupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL)
+                .build();
+
+        // All forms of amplitudes are valid when amplitude control is available.
+        assertTrue(info.areVibrationFeaturesSupported(
+                waveformWithAmplitudes(255, 0, VibrationEffect.DEFAULT_AMPLITUDE, 255)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createWaveform(
+                        /* timings= */ new long[] {1, 2, 3}, /* repeatIndex= */ -1)));
+        assertTrue(info.areVibrationFeaturesSupported(waveformWithAmplitudes(10, 30, 50)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                waveformWithAmplitudes(7, 255, 0, 0, 60)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createOneShot(20, VibrationEffect.DEFAULT_AMPLITUDE)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createOneShot(20, /* amplitude= */ 255)));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.createOneShot(20, /* amplitude= */ 40)));
+    }
+
+    @Test
+    public void areVibrationFeaturesSupported_compositionsWithSupportedPrimitivesSupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
+                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
+                .build();
+
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                        .compose()));
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(
+                                VibrationEffect.Composition.PRIMITIVE_CLICK,
+                                /* scale= */ 0.2f,
+                                /* delay= */ 200)
+                        .compose()));
+    }
+
+    @Test
+    public void areVibrationFeaturesSupported_compositionsWithUnupportedPrimitivesUnsupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
+                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
+                .build();
+
+        assertFalse(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD)
+                        .compose()));
+        assertFalse(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_LOW_TICK)
+                        .compose()));
+    }
+
+    @Test
+    public void areVibrationFeaturesSupported_composedEffects_allComponentsSupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS | IVibrator.CAP_AMPLITUDE_CONTROL)
+                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
+                .setSupportedEffects(VibrationEffect.EFFECT_TICK, VibrationEffect.EFFECT_POP)
+                .build();
+
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                        .addEffect(VibrationEffect.createWaveform(
+                                /* timings= */ new long[] {1, 2, 3},
+                                /* amplitudes= */ new int[] {10, 20, 255},
+                                /* repeatIndex= */ -1))
+                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_CLICK))
+                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_POP))
+                        .compose()));
+
+        info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
+                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD, 10)
+                .setSupportedEffects(VibrationEffect.EFFECT_POP, VibrationEffect.EFFECT_CLICK)
+                .build();
+
+        assertTrue(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD)
+                        .addEffect(VibrationEffect.createWaveform(
+                                // These timings are given either 0 or default amplitudes, which
+                                // do not require vibrator's amplitude control.
+                                /* timings= */ new long[] {1, 2, 3},
+                                /* repeatIndex= */ -1))
+                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_POP))
+                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_CLICK))
+                        .compose()));
+    }
+
+    @Test
+    public void areVibrationFeaturesSupported_composedEffects_someComponentsUnupported() {
+        VibratorInfo info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS | IVibrator.CAP_AMPLITUDE_CONTROL)
+                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
+                .setSupportedEffects(VibrationEffect.EFFECT_TICK, VibrationEffect.EFFECT_POP)
+                .build();
+
+        // Not supported due to the TICK primitive, which the vibrator has no support for.
+        assertFalse(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK)
+                        .addEffect(VibrationEffect.createWaveform(
+                                /* timings= */ new long[] {1, 2, 3},
+                                /* amplitudes= */ new int[] {10, 20, 255},
+                                /* repeatIndex= */ -1))
+                        .compose()));
+        // Not supported due to the THUD effect, which the vibrator has no support for.
+        assertFalse(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                        .addEffect(VibrationEffect.createWaveform(
+                                /* timings= */ new long[] {1, 2, 3},
+                                /* amplitudes= */ new int[] {10, 20, 255},
+                                /* repeatIndex= */ -1))
+                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_THUD))
+                        .compose()));
+
+        info = new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
+                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD, 10)
+                .setSupportedEffects(VibrationEffect.EFFECT_POP)
+                .build();
+
+        // Not supported due to fractional amplitudes (amplitudes not min (0) or max (255) or
+        // DEFAULT), because the vibrator has no amplitude control.
+        assertFalse(info.areVibrationFeaturesSupported(
+                VibrationEffect.startComposition()
+                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD)
+                        .addEffect(VibrationEffect.createWaveform(
+                                /* timings= */ new long[] {1, 2, 3},
+                                /* amplitudes= */ new int[] {10, 20, 255},
+                                /* repeatIndex= */ -1))
+                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_POP))
+                        .compose()));
+    }
+
+    private static VibrationEffect waveformWithAmplitudes(int...amplitudes) {
+        long[] timings = new long[amplitudes.length];
+        for (int i = 0; i < timings.length; i++) {
+            timings[i] = i * 2; // Arbitrary timings.
+        }
+        return VibrationEffect.createWaveform(timings, amplitudes, /* repeatIndex= */ -1);
+    }
 }
diff --git a/core/tests/vibrator/src/android/os/VibratorTest.java b/core/tests/vibrator/src/android/os/VibratorTest.java
index c559e34..8141ca4 100644
--- a/core/tests/vibrator/src/android/os/VibratorTest.java
+++ b/core/tests/vibrator/src/android/os/VibratorTest.java
@@ -578,189 +578,6 @@
         assertEquals(new VibrationAttributes.Builder().build(), vibrationAttributes);
     }
 
-    @Test
-    public void areVibrationFeaturesSupported_noAmplitudeControl_fractionalAmplitudes() {
-        Vibrator vibrator =
-                createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                        .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                        .setSupportedEffects(VibrationEffect.EFFECT_THUD)
-                        .build());
-
-        // Have at least one fractional amplitude (amplitude not min (0) or max (255) or DEFAULT).
-        assertFalse(vibrator.areVibrationFeaturesSupported(waveformWithAmplitudes(10, 30)));
-        assertFalse(vibrator.areVibrationFeaturesSupported(waveformWithAmplitudes(10, 255)));
-        assertFalse(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createOneShot(20, /* amplitude= */ 40)));
-    }
-
-    @Test
-    public void areVibrationFeaturesSupported_noAmplitudeControl_nonFractionalAmplitudes() {
-        Vibrator vibrator =
-                createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                        .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                        .setSupportedEffects(VibrationEffect.EFFECT_THUD)
-                        .build());
-
-        // All amplitudes are min, max, or default. Requires no amplitude control.
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                waveformWithAmplitudes(255, 0, VibrationEffect.DEFAULT_AMPLITUDE, 255)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createWaveform(
-                        /* timings= */ new long[] {1, 2, 3}, /* repeatIndex= */ -1)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createOneShot(20, VibrationEffect.DEFAULT_AMPLITUDE)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createOneShot(20, /* amplitude= */ 255)));
-    }
-
-    @Test
-    public void areVibrationFeaturesSupported_withAmplitudeControl() {
-        Vibrator vibrator =
-                createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                        .setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL)
-                        .build());
-
-        // All forms of amplitudes are valid when amplitude control is available.
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                waveformWithAmplitudes(255, 0, VibrationEffect.DEFAULT_AMPLITUDE, 255)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createWaveform(
-                        /* timings= */ new long[] {1, 2, 3}, /* repeatIndex= */ -1)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(waveformWithAmplitudes(10, 30, 50)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                waveformWithAmplitudes(7, 255, 0, 0, 60)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createOneShot(20, VibrationEffect.DEFAULT_AMPLITUDE)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createOneShot(20, /* amplitude= */ 255)));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.createOneShot(20, /* amplitude= */ 40)));
-    }
-
-    @Test
-    public void areVibrationFeaturesSupported_primitiveCompositionsWithSupportedPrimitives() {
-        Vibrator vibrator = createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
-                .build());
-
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
-                        .compose()));
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(
-                                VibrationEffect.Composition.PRIMITIVE_CLICK,
-                                /* scale= */ 0.2f,
-                                /* delay= */ 200)
-                        .compose()));
-    }
-
-    @Test
-    public void areVibrationFeaturesSupported_primitiveCompositionsWithUnupportedPrimitives() {
-        Vibrator vibrator = createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
-                .build());
-
-        assertFalse(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD)
-                        .compose()));
-        assertFalse(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_LOW_TICK)
-                        .compose()));
-    }
-
-    @Test
-    public void areVibrationFeaturesSupported_composedEffects_allComponentsSupported() {
-        Vibrator vibrator = createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS | IVibrator.CAP_AMPLITUDE_CONTROL)
-                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
-                .setSupportedEffects(VibrationEffect.EFFECT_TICK, VibrationEffect.EFFECT_POP)
-                .build());
-
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
-                        .addEffect(VibrationEffect.createWaveform(
-                                /* timings= */ new long[] {1, 2, 3},
-                                /* amplitudes= */ new int[] {10, 20, 255},
-                                /* repeatIndex= */ -1))
-                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_CLICK))
-                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_POP))
-                        .compose()));
-
-        vibrator = createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD, 10)
-                .setSupportedEffects(VibrationEffect.EFFECT_POP, VibrationEffect.EFFECT_CLICK)
-                .build());
-
-        assertTrue(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD)
-                        .addEffect(VibrationEffect.createWaveform(
-                                // These timings are given either 0 or default amplitudes, which
-                                // do not require vibrator's amplitude control.
-                                /* timings= */ new long[] {1, 2, 3},
-                                /* repeatIndex= */ -1))
-                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_POP))
-                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_CLICK))
-                        .compose()));
-    }
-
-    @Test
-    public void areVibrationFeaturesSupported_composedEffects_someComponentsUnupported() {
-        Vibrator vibrator = createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS | IVibrator.CAP_AMPLITUDE_CONTROL)
-                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 10)
-                .setSupportedEffects(VibrationEffect.EFFECT_TICK, VibrationEffect.EFFECT_POP)
-                .build());
-
-        // Not supported due to the TICK primitive, which the vibrator has no support for.
-        assertFalse(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK)
-                        .addEffect(VibrationEffect.createWaveform(
-                                /* timings= */ new long[] {1, 2, 3},
-                                /* amplitudes= */ new int[] {10, 20, 255},
-                                /* repeatIndex= */ -1))
-                        .compose()));
-        // Not supported due to the THUD effect, which the vibrator has no support for.
-        assertFalse(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
-                        .addEffect(VibrationEffect.createWaveform(
-                                /* timings= */ new long[] {1, 2, 3},
-                                /* amplitudes= */ new int[] {10, 20, 255},
-                                /* repeatIndex= */ -1))
-                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_THUD))
-                        .compose()));
-
-        vibrator = createVibratorWithCustomInfo(new VibratorInfo.Builder(/* id= */ 1)
-                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                .setSupportedPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD, 10)
-                .setSupportedEffects(VibrationEffect.EFFECT_POP)
-                .build());
-
-        // Not supported due to fractional amplitudes (amplitudes not min (0) or max (255) or
-        // DEFAULT), because the vibrator has no amplitude control.
-        assertFalse(vibrator.areVibrationFeaturesSupported(
-                VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_THUD)
-                        .addEffect(VibrationEffect.createWaveform(
-                                /* timings= */ new long[] {1, 2, 3},
-                                /* amplitudes= */ new int[] {10, 20, 255},
-                                /* repeatIndex= */ -1))
-                        .addEffect(VibrationEffect.createPredefined(VibrationEffect.EFFECT_POP))
-                        .compose()));
-    }
-
     /**
      * Asserts that the frequency profile is empty, and therefore frequency control isn't supported.
      */
@@ -768,21 +585,4 @@
         assertTrue(info.getFrequencyProfile().isEmpty());
         assertEquals(false, info.hasCapability(IVibrator.CAP_FREQUENCY_CONTROL));
     }
-
-    private Vibrator createVibratorWithCustomInfo(VibratorInfo info) {
-        return new SystemVibrator(mContextSpy) {
-            @Override
-            public VibratorInfo getInfo() {
-                return info;
-            }
-        };
-    }
-
-    private static VibrationEffect waveformWithAmplitudes(int...amplitudes) {
-        long[] timings = new long[amplitudes.length];
-        for (int i = 0; i < timings.length; i++) {
-            timings[i] = i * 2; // Arbitrary timings.
-        }
-        return VibrationEffect.createWaveform(timings, amplitudes, /* repeatIndex= */ -1);
-    }
 }
diff --git a/core/tests/vibrator/src/android/os/vibrator/PrebakedSegmentTest.java b/core/tests/vibrator/src/android/os/vibrator/PrebakedSegmentTest.java
index 8268077..4f5f3c0 100644
--- a/core/tests/vibrator/src/android/os/vibrator/PrebakedSegmentTest.java
+++ b/core/tests/vibrator/src/android/os/vibrator/PrebakedSegmentTest.java
@@ -25,18 +25,14 @@
 import static org.testng.Assert.assertThrows;
 
 import android.os.Parcel;
-import android.os.SystemVibrator;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
 import android.os.VibratorInfo;
 
-import androidx.test.InstrumentationRegistry;
-
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.junit.MockitoJUnitRunner;
+import org.junit.runners.JUnit4;
 
-@RunWith(MockitoJUnitRunner.class)
+@RunWith(JUnit4.class)
 public class PrebakedSegmentTest {
 
     @Test
@@ -149,121 +145,121 @@
 
     @Test
     public void testVibrationFeaturesSupport_idsWithFallback_fallbackEnabled_vibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(
                 VibrationEffect.EFFECT_TICK,
                 VibrationEffect.EFFECT_CLICK,
                 VibrationEffect.EFFECT_DOUBLE_CLICK,
                 VibrationEffect.EFFECT_HEAVY_CLICK);
 
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_DOUBLE_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_HEAVY_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
 
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithFallback_fallbackEnabled_noVibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(new int[0]);
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(new int[0]);
 
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_DOUBLE_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_HEAVY_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithFallback_fallbackDisabled_vibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(
                 VibrationEffect.EFFECT_TICK,
                 VibrationEffect.EFFECT_CLICK,
                 VibrationEffect.EFFECT_DOUBLE_CLICK,
                 VibrationEffect.EFFECT_HEAVY_CLICK);
 
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_DOUBLE_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_HEAVY_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithFallback_fallbackDisabled_noVibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(new int[0]);
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(new int[0]);
 
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_DOUBLE_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_HEAVY_CLICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithNoFallback_fallbackEnabled_vibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(
                 VibrationEffect.EFFECT_THUD,
                 VibrationEffect.EFFECT_POP,
                 VibrationEffect.EFFECT_TEXTURE_TICK);
 
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_THUD)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_POP)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithFallback(VibrationEffect.EFFECT_TEXTURE_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithNoFallback_fallbackEnabled_noVibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(new int[0]);
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(new int[0]);
 
         assertFalse(createSegmentWithFallback(VibrationEffect.EFFECT_THUD)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithFallback(VibrationEffect.EFFECT_POP)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithFallback(VibrationEffect.EFFECT_TEXTURE_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithNoFallback_fallbackDisabled_vibratorSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(
                 VibrationEffect.EFFECT_THUD,
                 VibrationEffect.EFFECT_POP,
                 VibrationEffect.EFFECT_TEXTURE_TICK);
 
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_THUD)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_POP)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertTrue(createSegmentWithoutFallback(VibrationEffect.EFFECT_TEXTURE_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_idsWithNoFallback_fallbackDisabled_noVibSupport() {
-        Vibrator vibrator = createVibratorWithSupportedEffects(new int[0]);
+        VibratorInfo info = createVibratorInfoWithSupportedEffects(new int[0]);
 
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_THUD)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_POP)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
         assertFalse(createSegmentWithoutFallback(VibrationEffect.EFFECT_TEXTURE_TICK)
-                .areVibrationFeaturesSupported(vibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
@@ -283,14 +279,9 @@
         return new PrebakedSegment(effectId, false, VibrationEffect.EFFECT_STRENGTH_MEDIUM);
     }
 
-    private static Vibrator createVibratorWithSupportedEffects(int... supportedEffects) {
-        return new SystemVibrator(InstrumentationRegistry.getContext()) {
-            @Override
-            public VibratorInfo getInfo() {
-                return new VibratorInfo.Builder(/* id= */ 1)
-                        .setSupportedEffects(supportedEffects)
-                        .build();
-            }
-        };
+    private static VibratorInfo createVibratorInfoWithSupportedEffects(int... supportedEffects) {
+        return new VibratorInfo.Builder(/* id= */ 1)
+                .setSupportedEffects(supportedEffects)
+                .build();
     }
 }
diff --git a/core/tests/vibrator/src/android/os/vibrator/PrimitiveSegmentTest.java b/core/tests/vibrator/src/android/os/vibrator/PrimitiveSegmentTest.java
index 6f5adcd..ec5a084 100644
--- a/core/tests/vibrator/src/android/os/vibrator/PrimitiveSegmentTest.java
+++ b/core/tests/vibrator/src/android/os/vibrator/PrimitiveSegmentTest.java
@@ -25,18 +25,14 @@
 
 import android.hardware.vibrator.IVibrator;
 import android.os.Parcel;
-import android.os.SystemVibrator;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
 import android.os.VibratorInfo;
 
-import androidx.test.InstrumentationRegistry;
-
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.junit.MockitoJUnitRunner;
+import org.junit.runners.JUnit4;
 
-@RunWith(MockitoJUnitRunner.class)
+@RunWith(JUnit4.class)
 public class PrimitiveSegmentTest {
     private static final float TOLERANCE = 1e-2f;
 
@@ -146,15 +142,15 @@
     public void testVibrationFeaturesSupport_primitiveSupportedByVibrator() {
         assertTrue(createSegment(VibrationEffect.Composition.PRIMITIVE_CLICK)
                 .areVibrationFeaturesSupported(
-                        createVibratorWithSupportedPrimitive(
+                        createVibratorInfoWithSupportedPrimitive(
                                 VibrationEffect.Composition.PRIMITIVE_CLICK)));
         assertTrue(createSegment(VibrationEffect.Composition.PRIMITIVE_THUD)
                 .areVibrationFeaturesSupported(
-                        createVibratorWithSupportedPrimitive(
+                        createVibratorInfoWithSupportedPrimitive(
                                 VibrationEffect.Composition.PRIMITIVE_THUD)));
         assertTrue(createSegment(VibrationEffect.Composition.PRIMITIVE_QUICK_RISE)
                 .areVibrationFeaturesSupported(
-                        createVibratorWithSupportedPrimitive(
+                        createVibratorInfoWithSupportedPrimitive(
                                 VibrationEffect.Composition.PRIMITIVE_QUICK_RISE)));
     }
 
@@ -162,15 +158,15 @@
     public void testVibrationFeaturesSupport_primitiveNotSupportedByVibrator() {
         assertFalse(createSegment(VibrationEffect.Composition.PRIMITIVE_CLICK)
                 .areVibrationFeaturesSupported(
-                        createVibratorWithSupportedPrimitive(
+                        createVibratorInfoWithSupportedPrimitive(
                                 VibrationEffect.Composition.PRIMITIVE_THUD)));
         assertFalse(createSegment(VibrationEffect.Composition.PRIMITIVE_THUD)
                 .areVibrationFeaturesSupported(
-                        createVibratorWithSupportedPrimitive(
+                        createVibratorInfoWithSupportedPrimitive(
                                 VibrationEffect.Composition.PRIMITIVE_CLICK)));
         assertFalse(createSegment(VibrationEffect.Composition.PRIMITIVE_THUD)
                 .areVibrationFeaturesSupported(
-                        createVibratorWithSupportedPrimitive(
+                        createVibratorInfoWithSupportedPrimitive(
                                 VibrationEffect.Composition.PRIMITIVE_QUICK_RISE)));
     }
 
@@ -193,15 +189,10 @@
         return new PrimitiveSegment(primitiveId, 0.2f, 10);
     }
 
-    private static Vibrator createVibratorWithSupportedPrimitive(int primitiveId) {
-        return new SystemVibrator(InstrumentationRegistry.getContext()) {
-            @Override
-            public VibratorInfo getInfo() {
-                return new VibratorInfo.Builder(/* id= */ 1)
-                        .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
-                        .setSupportedPrimitive(primitiveId, 10)
-                        .build();
-                }
-        };
+    private static VibratorInfo createVibratorInfoWithSupportedPrimitive(int primitiveId) {
+        return new VibratorInfo.Builder(/* id= */ 1)
+                .setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS)
+                .setSupportedPrimitive(primitiveId, 10)
+                .build();
     }
 }
diff --git a/core/tests/vibrator/src/android/os/vibrator/RampSegmentTest.java b/core/tests/vibrator/src/android/os/vibrator/RampSegmentTest.java
index 68870e5..5caa86b 100644
--- a/core/tests/vibrator/src/android/os/vibrator/RampSegmentTest.java
+++ b/core/tests/vibrator/src/android/os/vibrator/RampSegmentTest.java
@@ -23,31 +23,21 @@
 import static junit.framework.Assert.assertSame;
 import static junit.framework.Assert.assertTrue;
 
-import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertThrows;
 
+import android.hardware.vibrator.IVibrator;
 import android.os.Parcel;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
+import android.os.VibratorInfo;
 
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnit;
-import org.mockito.junit.MockitoJUnitRunner;
-import org.mockito.junit.MockitoRule;
+import org.junit.runners.JUnit4;
 
-@RunWith(MockitoJUnitRunner.class)
+@RunWith(JUnit4.class)
 public class RampSegmentTest {
     private static final float TOLERANCE = 1e-2f;
 
-    @Rule
-    public MockitoRule mMockitoRule = MockitoJUnit.rule();
-
-    @Mock
-    private Vibrator mVibrator;
-
     @Test
     public void testCreation() {
         RampSegment ramp = new RampSegment(/* startAmplitude= */ 1, /* endAmplitude= */ 0,
@@ -147,71 +137,71 @@
 
     @Test
     public void testVibrationFeaturesSupport_amplitudeAndFrequencyControls_supported() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ true);
 
         // Increasing amplitude
-        assertTrue(new RampSegment(0.5f, 1, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
+        assertTrue(new RampSegment(0.5f, 1, 0, 0, 10).areVibrationFeaturesSupported(info));
         // Increasing frequency
-        assertTrue(new RampSegment(0.5f, 0.5f, 0, 1, 10).areVibrationFeaturesSupported(mVibrator));
+        assertTrue(new RampSegment(0.5f, 0.5f, 0, 1, 10).areVibrationFeaturesSupported(info));
         // Decreasing amplitude
-        assertTrue(new RampSegment(1, 0.5f, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
+        assertTrue(new RampSegment(1, 0.5f, 0, 0, 10).areVibrationFeaturesSupported(info));
         // Decreasing frequency
-        assertTrue(new RampSegment(0.5f, 0.5f, 1, 0, 10).areVibrationFeaturesSupported(mVibrator));
+        assertTrue(new RampSegment(0.5f, 0.5f, 1, 0, 10).areVibrationFeaturesSupported(info));
         // Zero duration
-        assertTrue(new RampSegment(0.5f, 0.5f, 1, 0, 0).areVibrationFeaturesSupported(mVibrator));
+        assertTrue(new RampSegment(0.5f, 0.5f, 1, 0, 0).areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_noAmplitudeControl_unsupportedForChangingAmplitude() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ false, /* hasFrequencyControl= */ true);
 
         // Test with increasing/decreasing amplitudes.
-        assertFalse(new RampSegment(0.5f, 1, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
-        assertFalse(new RampSegment(1, 0.5f, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
+        assertFalse(new RampSegment(0.5f, 1, 0, 0, 10).areVibrationFeaturesSupported(info));
+        assertFalse(new RampSegment(1, 0.5f, 0, 0, 10).areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_noAmplitudeControl_fractionalAmplitudeUnsupported() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ false, /* hasFrequencyControl= */ true);
 
-        assertFalse(new RampSegment(0.2f, 0.2f, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
-        assertFalse(new RampSegment(0, 0.2f, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
-        assertFalse(new RampSegment(0.2f, 0, 0, 0, 10).areVibrationFeaturesSupported(mVibrator));
+        assertFalse(new RampSegment(0.2f, 0.2f, 0, 0, 10).areVibrationFeaturesSupported(info));
+        assertFalse(new RampSegment(0, 0.2f, 0, 0, 10).areVibrationFeaturesSupported(info));
+        assertFalse(new RampSegment(0.2f, 0, 0, 0, 10).areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_unchangingZeroAmplitude_supported() {
         RampSegment amplitudeZeroWithIncreasingFrequency = new RampSegment(1, 1, 0.5f, 0.8f, 10);
         RampSegment amplitudeZeroWithDecreasingFrequency = new RampSegment(1, 1, 0.8f, 0.5f, 10);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ false, /* hasFrequencyControl= */ true);
 
-        assertTrue(amplitudeZeroWithIncreasingFrequency.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(amplitudeZeroWithDecreasingFrequency.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(amplitudeZeroWithIncreasingFrequency.areVibrationFeaturesSupported(info));
+        assertTrue(amplitudeZeroWithDecreasingFrequency.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        info = createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ true);
 
-        assertTrue(amplitudeZeroWithIncreasingFrequency.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(amplitudeZeroWithDecreasingFrequency.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(amplitudeZeroWithIncreasingFrequency.areVibrationFeaturesSupported(info));
+        assertTrue(amplitudeZeroWithDecreasingFrequency.areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_unchangingOneAmplitude_supported() {
         RampSegment amplitudeOneWithIncreasingFrequency = new RampSegment(1, 1, 0.5f, 0.8f, 10);
         RampSegment amplitudeOneWithDecreasingFrequency = new RampSegment(1, 1, 0.8f, 0.5f, 10);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ false, /* hasFrequencyControl= */ true);
 
-        assertTrue(amplitudeOneWithIncreasingFrequency.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(amplitudeOneWithDecreasingFrequency.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(amplitudeOneWithIncreasingFrequency.areVibrationFeaturesSupported(info));
+        assertTrue(amplitudeOneWithDecreasingFrequency.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        info = createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ true);
 
-        assertTrue(amplitudeOneWithIncreasingFrequency.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(amplitudeOneWithDecreasingFrequency.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(amplitudeOneWithIncreasingFrequency.areVibrationFeaturesSupported(info));
+        assertTrue(amplitudeOneWithDecreasingFrequency.areVibrationFeaturesSupported(info));
     }
 
     @Test
@@ -220,52 +210,52 @@
                 new RampSegment(DEFAULT_AMPLITUDE, DEFAULT_AMPLITUDE, 0.5f, 0.8f, 10);
         RampSegment defaultAmplitudeDecreasingFrequency =
                 new RampSegment(DEFAULT_AMPLITUDE, DEFAULT_AMPLITUDE, 0.8f, 0.5f, 10);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ false, /* hasFrequencyControl= */ true);
 
-        assertTrue(defaultAmplitudeIncreasingFrequency.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(defaultAmplitudeDecreasingFrequency.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(defaultAmplitudeIncreasingFrequency.areVibrationFeaturesSupported(info));
+        assertTrue(defaultAmplitudeDecreasingFrequency.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        info = createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ true);
 
-        assertTrue(defaultAmplitudeIncreasingFrequency.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(defaultAmplitudeDecreasingFrequency.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(defaultAmplitudeIncreasingFrequency.areVibrationFeaturesSupported(info));
+        assertTrue(defaultAmplitudeDecreasingFrequency.areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_noFrequencyControl_unsupportedForChangingFrequency() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
-        when(mVibrator.hasFrequencyControl()).thenReturn(false);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ false);
 
         // Test with increasing/decreasing frequencies.
-        assertFalse(new RampSegment(0, 0, 0.2f, 0.4f, 10).areVibrationFeaturesSupported(mVibrator));
-        assertFalse(new RampSegment(0, 0, 0.4f, 0.2f, 10).areVibrationFeaturesSupported(mVibrator));
+        assertFalse(new RampSegment(0, 0, 0.2f, 0.4f, 10).areVibrationFeaturesSupported(info));
+        assertFalse(new RampSegment(0, 0, 0.4f, 0.2f, 10).areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_noFrequencyControl_fractionalFrequencyUnsupported() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
-        when(mVibrator.hasFrequencyControl()).thenReturn(false);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ false);
 
-        assertFalse(new RampSegment(0, 0, 0.2f, 0.2f, 10).areVibrationFeaturesSupported(mVibrator));
-        assertFalse(new RampSegment(0, 0, 0.2f, 0, 10).areVibrationFeaturesSupported(mVibrator));
-        assertFalse(new RampSegment(0, 0, 0, 0.2f, 10).areVibrationFeaturesSupported(mVibrator));
+        assertFalse(new RampSegment(0, 0, 0.2f, 0.2f, 10).areVibrationFeaturesSupported(info));
+        assertFalse(new RampSegment(0, 0, 0.2f, 0, 10).areVibrationFeaturesSupported(info));
+        assertFalse(new RampSegment(0, 0, 0, 0.2f, 10).areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_unchangingZeroFrequency_supported() {
         RampSegment frequencyZeroWithIncreasingAmplitude = new RampSegment(0.1f, 1, 0, 0, 10);
         RampSegment frequencyZeroWithDecreasingAmplitude = new RampSegment(1, 0.1f, 0, 0, 10);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
-        when(mVibrator.hasFrequencyControl()).thenReturn(false);
+        VibratorInfo info =
+                createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ false);
 
-        assertTrue(frequencyZeroWithIncreasingAmplitude.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(frequencyZeroWithDecreasingAmplitude.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(frequencyZeroWithIncreasingAmplitude.areVibrationFeaturesSupported(info));
+        assertTrue(frequencyZeroWithDecreasingAmplitude.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
+        info = createVibInfo(/* hasAmplitudeControl= */ true, /* hasFrequencyControl= */ true);
 
-        assertTrue(frequencyZeroWithIncreasingAmplitude.areVibrationFeaturesSupported(mVibrator));
-        assertTrue(frequencyZeroWithDecreasingAmplitude.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(frequencyZeroWithIncreasingAmplitude.areVibrationFeaturesSupported(info));
+        assertTrue(frequencyZeroWithDecreasingAmplitude.areVibrationFeaturesSupported(info));
     }
 
     @Test
@@ -274,4 +264,17 @@
         // duration checked in VibrationEffect implementations.
         assertTrue(new RampSegment(0.5f, 1, 0, 0, 5_000).isHapticFeedbackCandidate());
     }
+
+    private static VibratorInfo createVibInfo(
+            boolean hasAmplitudeControl, boolean hasFrequencyControl) {
+        VibratorInfo.Builder builder = new VibratorInfo.Builder(/* id= */ 1);
+        long capabilities = 0;
+        if (hasAmplitudeControl) {
+            capabilities |= IVibrator.CAP_AMPLITUDE_CONTROL;
+        }
+        if (hasFrequencyControl) {
+            capabilities |= (IVibrator.CAP_FREQUENCY_CONTROL | IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
+        }
+        return builder.setCapabilities(capabilities).build();
+    }
 }
diff --git a/core/tests/vibrator/src/android/os/vibrator/StepSegmentTest.java b/core/tests/vibrator/src/android/os/vibrator/StepSegmentTest.java
index 34bb892..44db306 100644
--- a/core/tests/vibrator/src/android/os/vibrator/StepSegmentTest.java
+++ b/core/tests/vibrator/src/android/os/vibrator/StepSegmentTest.java
@@ -21,31 +21,20 @@
 import static junit.framework.Assert.assertSame;
 import static junit.framework.Assert.assertTrue;
 
-import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertThrows;
 
+import android.hardware.vibrator.IVibrator;
 import android.os.Parcel;
 import android.os.VibrationEffect;
-import android.os.Vibrator;
+import android.os.VibratorInfo;
 
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnit;
-import org.mockito.junit.MockitoJUnitRunner;
-import org.mockito.junit.MockitoRule;
+import org.junit.runners.JUnit4;
 
-@RunWith(MockitoJUnitRunner.class)
+@RunWith(JUnit4.class)
 public class StepSegmentTest {
     private static final float TOLERANCE = 1e-2f;
-
-    @Rule
-    public MockitoRule mMockitoRule = MockitoJUnit.rule();
-
-    @Mock
-    private Vibrator mVibrator;
-
     @Test
     public void testCreation() {
         StepSegment step = new StepSegment(/* amplitude= */ 1f, /* frequencyHz= */ 1f,
@@ -160,26 +149,26 @@
     public void testVibrationFeaturesSupport_zeroAmplitude_supported() {
         StepSegment segment =
                 new StepSegment(/* amplitude= */ 0, /* frequencyHz= */ 0, /* duration= */ 0);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        VibratorInfo info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ true);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ false);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_maxAmplitude_supported() {
         StepSegment segment =
                 new StepSegment(/* amplitude= */ 1, /* frequencyHz= */ 0, /* duration= */ 0);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        VibratorInfo info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ true);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ false);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
     }
 
     @Test
@@ -189,60 +178,60 @@
                         /* amplitude= */ VibrationEffect.DEFAULT_AMPLITUDE,
                         /* frequencyHz= */ 0,
                         /* duration= */ 0);
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        VibratorInfo info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ true);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ false);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_fractionalAmplitude_hasAmplitudeCtrl_supported() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(true);
+        VibratorInfo info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ true);
 
         assertTrue(new StepSegment(/* amplitude= */ 0.2f, /* frequencyHz= */ 0, /* duration= */ 0)
-                .areVibrationFeaturesSupported(mVibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_fractionalAmplitude_hasNoAmplitudeCtrl_notSupported() {
-        when(mVibrator.hasAmplitudeControl()).thenReturn(false);
+        VibratorInfo info = createVibInfoForAmplitude(/* hasAmplitudeControl= */ false);
 
         assertFalse(new StepSegment(/* amplitude= */ 0.2f, /* frequencyHz= */ 0, /* duration= */ 0)
-                .areVibrationFeaturesSupported(mVibrator));
+                .areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_zeroFrequency_supported() {
         StepSegment segment =
                 new StepSegment(/* amplitude= */ 0f, /* frequencyHz= */ 0, /* duration= */ 0);
-        when(mVibrator.hasFrequencyControl()).thenReturn(false);
+        VibratorInfo info = createVibInfoForFrequency(/* hasFrequencyControl= */ false);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
 
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
+        info = createVibInfoForFrequency(/* hasFrequencyControl= */ true);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_nonZeroFrequency_hasFrequencyCtrl_supported() {
         StepSegment segment =
                 new StepSegment(/* amplitude= */ 0f, /* frequencyHz= */ 0.2f, /* duration= */ 0);
-        when(mVibrator.hasFrequencyControl()).thenReturn(true);
+        VibratorInfo info = createVibInfoForFrequency(/* hasFrequencyControl= */ true);
 
-        assertTrue(segment.areVibrationFeaturesSupported(mVibrator));
+        assertTrue(segment.areVibrationFeaturesSupported(info));
     }
 
     @Test
     public void testVibrationFeaturesSupport_nonZeroFrequency_hasNoFrequencyCtrl_notSupported() {
         StepSegment segment =
                 new StepSegment(/* amplitude= */ 0f, /* frequencyHz= */ 0.2f, /* duration= */ 0);
-        when(mVibrator.hasFrequencyControl()).thenReturn(false);
+        VibratorInfo info = createVibInfoForFrequency(/* hasFrequencyControl= */ false);
 
-        assertFalse(segment.areVibrationFeaturesSupported(mVibrator));
+        assertFalse(segment.areVibrationFeaturesSupported(info));
     }
 
     @Test
@@ -251,4 +240,21 @@
         // duration checked in VibrationEffect implementations.
         assertTrue(new StepSegment(0, 0, 5_000).isHapticFeedbackCandidate());
     }
+
+    private static VibratorInfo createVibInfoForAmplitude(boolean hasAmplitudeControl) {
+        VibratorInfo.Builder builder = new VibratorInfo.Builder(/* id= */ 1);
+        if (hasAmplitudeControl) {
+            builder.setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL);
+        }
+        return builder.build();
+    }
+
+    private static VibratorInfo createVibInfoForFrequency(boolean hasFrequencyControl) {
+        VibratorInfo.Builder builder = new VibratorInfo.Builder(/* id= */ 1);
+        if (hasFrequencyControl) {
+            builder.setCapabilities(
+                    IVibrator.CAP_FREQUENCY_CONTROL | IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
+        }
+        return builder.build();
+    }
 }