/*
 * Copyright 2022 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.media3.transformer.mh.analysis;

import static androidx.media3.common.util.Assertions.checkNotNull;
import static com.google.common.truth.Truth.assertThat;

import android.media.MediaFormat;
import androidx.annotation.Nullable;
import androidx.media3.common.C;
import androidx.media3.common.ColorInfo;
import androidx.media3.common.util.MediaFormatUtil;
import androidx.media3.common.util.Util;
import androidx.media3.test.utils.DecodeOneFrameUtil;

/** Utilities for reading color info from a file. */
public class FileUtil {
  public static void assertFileHasColorTransfer(
      @Nullable String filePath, @C.ColorTransfer int expectedColorTransfer) throws Exception {
    if (Util.SDK_INT < 29) {
      // Skipping on this API version due to lack of support for MediaFormat#getInteger, which is
      // required for MediaFormatUtil#getColorInfo.
      return;
    }
    DecodeOneFrameUtil.decodeOneCacheFileFrame(
        checkNotNull(filePath),
        new DecodeOneFrameUtil.Listener() {
          @Override
          public void onContainerExtracted(MediaFormat mediaFormat) {
            @Nullable ColorInfo extractedColor = MediaFormatUtil.getColorInfo(mediaFormat);
            assertThat(checkNotNull(extractedColor).colorTransfer).isEqualTo(expectedColorTransfer);
          }

          @Override
          public void onFrameDecoded(MediaFormat mediaFormat) {
            @Nullable ColorInfo decodedColor = MediaFormatUtil.getColorInfo(mediaFormat);
            assertThat(checkNotNull(decodedColor).colorTransfer).isEqualTo(expectedColorTransfer);
          }
        },
        /* surface= */ null);
  }

  private FileUtil() {}
}

@@ -1754,24 +1754,28 @@
 tractedColor
+Info
  = MediaForm
@@ -1830,67 +1830,49 @@
 sert
-That(checkNotNull(extractedColor).colorTransfer).isEqualTo(
+ColorInfoHasTransfer(extractedColorInfo, 
 expe
@@ -2032,16 +2032,20 @@
 dedColor
+Info
  = Media
@@ -2104,65 +2104,47 @@
 sert
-That(checkNotNull(decodedColor).colorTransfer).isEqualTo(
+ColorInfoHasTransfer(decodedColorInfo, 
 expe
@@ -2221,16 +2221,343 @@
 );%0A  %7D%0A%0A
+  private static void assertColorInfoHasTransfer(%0A      @Nullable ColorInfo colorInfo, @C.ColorTransfer int expectedColorTransfer) %7B%0A    @C.ColorTransfer%0A    int actualColorTransfer = colorInfo == null ? C.COLOR_TRANSFER_SDR : colorInfo.colorTransfer;%0A    assertThat(actualColorTransfer).isEqualTo(expectedColorTransfer);%0A  %7D%0A%0A
   privat

Update test FileUtil to handle null ColorInfo.