diff --git a/testutils/src/main/java/com/google/android/exoplayer2/testutil/Action.java b/testutils/src/main/java/com/google/android/exoplayer2/testutil/Action.java
index 835100101c..8fe58aa45b 100644
--- a/testutils/src/main/java/com/google/android/exoplayer2/testutil/Action.java
+++ b/testutils/src/main/java/com/google/android/exoplayer2/testutil/Action.java
@@ -890,13 +890,13 @@ public abstract class Action {
    */
   public static final class WaitForPlaybackState extends Action {
 
-    private final int targetPlaybackState;
+    @Player.State private final int targetPlaybackState;
 
     /**
      * @param tag A tag to use for logging.
      * @param targetPlaybackState The playback state to wait for.
      */
-    public WaitForPlaybackState(String tag, int targetPlaybackState) {
+    public WaitForPlaybackState(String tag, @Player.State int targetPlaybackState) {
       super(tag, "WaitForPlaybackState");
       this.targetPlaybackState = targetPlaybackState;
     }
diff --git a/testutils/src/main/java/com/google/android/exoplayer2/testutil/ActionSchedule.java b/testutils/src/main/java/com/google/android/exoplayer2/testutil/ActionSchedule.java
index d116fdc6b5..8051e997b3 100644
--- a/testutils/src/main/java/com/google/android/exoplayer2/testutil/ActionSchedule.java
+++ b/testutils/src/main/java/com/google/android/exoplayer2/testutil/ActionSchedule.java
@@ -537,7 +537,7 @@ public final class ActionSchedule {
      * @param targetPlaybackState The target playback state.
      * @return The builder, for convenience.
      */
-    public Builder waitForPlaybackState(int targetPlaybackState) {
+    public Builder waitForPlaybackState(@Player.State int targetPlaybackState) {
       return apply(new WaitForPlaybackState(tag, targetPlaybackState));
     }