diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/ts/AdtsReader.java b/libraries/extractor/src/main/java/androidx/media3/extractor/ts/AdtsReader.java
index 52543a11e8d..8f1a909f628 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/ts/AdtsReader.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/ts/AdtsReader.java
@@ -15,6 +15,7 @@
  */
 package androidx.media3.extractor.ts;
 
+import static androidx.media3.common.util.Assertions.checkNotNull;
 import static java.lang.Math.min;
 
 import androidx.annotation.Nullable;
@@ -33,6 +34,7 @@
 import androidx.media3.extractor.ExtractorOutput;
 import androidx.media3.extractor.TrackOutput;
 import androidx.media3.extractor.ts.TsPayloadReader.TrackIdGenerator;
+import com.google.common.collect.ImmutableList;
 import java.util.Arrays;
 import java.util.Collections;
 import org.checkerframework.checker.nullness.qual.EnsuresNonNull;
@@ -191,6 +193,7 @@ public void consume(ParsableByteArray data) throws ParserException {
           }
           break;
         case STATE_READING_AAC_PCE:
+          checkNotNull(pceBuffer);
           if (continueRead(data, pceBuffer.data, pceBuffer.data.length)) {
             readAacProgramConfigElement();
           }
@@ -571,8 +574,11 @@ private void parseAdtsHeader() throws ParserException {
     }
   }
 
-  @RequiresNonNull({"pendingOutputFormat", "pceBuffer"})
+  @RequiresNonNull("currentOutput")
   void readAacProgramConfigElement() {
+    Format pendingOutputFormat = checkNotNull(this.pendingOutputFormat);
+    ParsableBitArray pceBuffer = checkNotNull(this.pceBuffer);
+
     // See ISO 13818-7 Advanced Audio Coding (2006) Table 36 for PCE tag encoding.
     if (pceBuffer.readBits(3) == 5 /* PCE tag */) {
       // See ISO 13818-7 Advanced Audio Coding (2006) Table 25 for syntax of a PCE.
@@ -624,16 +630,15 @@ void readAacProgramConfigElement() {
 
           int configSize = oldConfig.length;
           configSize += (numPceBits + 7) / 8 + 1; // Byte align and add a zero length comment.
-          byte[] newConfig = new byte[configSize];
+          byte[] newConfig = Arrays.copyOf(oldConfig, configSize);
 
-          System.arraycopy(oldConfig, 0, newConfig, 0, oldConfig.length);
           pceBuffer.setPosition(3 /* PCE tag */);
           pceBuffer.readBits(newConfig, oldConfig.length, numPceBits);
 
           pendingOutputFormat =
               pendingOutputFormat
                   .buildUpon()
-                  .setInitializationData(Collections.singletonList(newConfig))
+                  .setInitializationData(ImmutableList.of(newConfig))
                   .build();
 
           // Submit PCE-appended output format.
@@ -643,13 +648,13 @@ void readAacProgramConfigElement() {
       }
     }
 
-    pendingOutputFormat = null;
-
     // Pass through all accumulated data as sample data.
     ParsableByteArray data = new ParsableByteArray(pceBuffer.data);
-    pceBuffer = null;
     setReadingSampleState(currentOutput, currentSampleDuration, 0, sampleSize);
     readSample(data);
+
+    this.pendingOutputFormat = null;
+    this.pceBuffer = null;
   }
 
   /** Reads the rest of the sample */
@@ -669,7 +674,7 @@ private void readSample(ParsableByteArray data) {
 
   @EnsuresNonNull({"output", "currentOutput", "id3Output"})
   private void assertTracksCreated() {
-    Assertions.checkNotNull(output);
+    checkNotNull(output);
     Util.castNonNull(currentOutput);
     Util.castNonNull(id3Output);
   }
diff --git a/libraries/extractor/src/test/java/androidx/media3/extractor/ts/AdtsReaderTest.java b/libraries/extractor/src/test/java/androidx/media3/extractor/ts/AdtsReaderTest.java
index 617236523d4..09ec2684a4c 100644
--- a/libraries/extractor/src/test/java/androidx/media3/extractor/ts/AdtsReaderTest.java
+++ b/libraries/extractor/src/test/java/androidx/media3/extractor/ts/AdtsReaderTest.java
@@ -17,6 +17,7 @@
 
 import static androidx.media3.extractor.ts.TsPayloadReader.FLAG_DATA_ALIGNMENT_INDICATOR;
 import static java.lang.Math.min;
+import static org.junit.Assert.assertThrows;
 
 import androidx.media3.common.C;
 import androidx.media3.common.ParserException;
@@ -208,7 +209,7 @@ public void aacPceData() throws ParserException {
     adtsOutput.assertSample(0, AAC_PCE_ADTS_CONTENT, 0, C.BUFFER_FLAG_KEY_FRAME, null);
   }
 
-  @Test(expected = IllegalStateException.class)
+  @Test
   public void aacPceDataFail() throws ParserException {
     data = new ParsableByteArray(Arrays.copyOf(AAC_PCE_TEST_DATA, AAC_PCE_TEST_DATA.length));
     byte[] bytes = data.getData();
@@ -218,7 +219,7 @@ public void aacPceDataFail() throws ParserException {
     bytes[AAC_PCE_ADTS_HEADER.length] |= 0x20;
 
     // Should throw as FakeTrackOutput expects a format before sampleMetadata.
-    feed();
+    assertThrows(IllegalStateException.class, this::feed);
   }
 
   private void feedLimited(int limit) throws ParserException {