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 b65accdf3f..2d6beff416 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
@@ -685,13 +685,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 f6ab4b9924..4800df662c 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
@@ -433,7 +433,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));
     }