diff --git a/media/tests/benchmark/MediaBenchmarkTest/src/androidTest/java/com/android/media/benchmark/tests/EncoderTest.java b/media/tests/benchmark/MediaBenchmarkTest/src/androidTest/java/com/android/media/benchmark/tests/EncoderTest.java
index 831467a..a5eda26 100644
--- a/media/tests/benchmark/MediaBenchmarkTest/src/androidTest/java/com/android/media/benchmark/tests/EncoderTest.java
+++ b/media/tests/benchmark/MediaBenchmarkTest/src/androidTest/java/com/android/media/benchmark/tests/EncoderTest.java
@@ -19,6 +19,9 @@
 import android.content.Context;
 import android.media.MediaCodec;
 import android.media.MediaFormat;
+
+import static android.media.MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible;
+
 import android.util.Log;
 
 import androidx.test.platform.app.InstrumentationRegistry;
@@ -65,6 +68,7 @@
     private static final int ENCODE_DEFAULT_FRAME_RATE = 25;
     private static final int ENCODE_DEFAULT_BIT_RATE = 8000000 /* 8 Mbps */;
     private static final int ENCODE_MIN_BIT_RATE = 600000 /* 600 Kbps */;
+    private static final int ENCODE_DEFAULT_AUDIO_BIT_RATE = 128000 /* 128 Kbps */;
     private String mInputFile;
 
     @Parameterized.Parameters
@@ -107,6 +111,7 @@
         int profile = 0;
         int level = 0;
         int frameRate = 0;
+
         //Parameters for audio
         int bitRate = 0;
         int sampleRate = 0;
@@ -122,6 +127,7 @@
         ArrayList<ByteBuffer> inputBuffer = new ArrayList<>();
         ArrayList<MediaCodec.BufferInfo> frameInfo = new ArrayList<>();
         for (int currentTrack = 0; currentTrack < trackCount; currentTrack++) {
+            int colorFormat = COLOR_FormatYUV420Flexible;
             extractor.selectExtractorTrack(currentTrack);
             MediaFormat format = extractor.getFormat(currentTrack);
             // Get samples from extractor
@@ -148,6 +154,7 @@
             status = decoder.decode(inputBuffer, frameInfo, false, format, "");
             assertEquals("Decoder returned error " + status + " for file: " + mInputFile, 0,
                     status);
+            MediaFormat decoderFormat = decoder.getFormat();
             decoder.deInitCodec();
             extractor.unselectExtractorTrack(currentTrack);
             inputBuffer.clear();
@@ -203,10 +210,17 @@
                         if (format.containsKey(MediaFormat.KEY_PROFILE)) {
                             level = format.getInteger(MediaFormat.KEY_LEVEL);
                         }
+                        if (decoderFormat.containsKey(MediaFormat.KEY_COLOR_FORMAT)) {
+                            colorFormat = decoderFormat.getInteger(MediaFormat.KEY_COLOR_FORMAT);
+                        }
                     } else {
                         sampleRate = format.getInteger(MediaFormat.KEY_SAMPLE_RATE);
                         numChannels = format.getInteger(MediaFormat.KEY_CHANNEL_COUNT);
-                        bitRate = sampleRate * numChannels * 16;
+                        if (decoderFormat.containsKey(MediaFormat.KEY_BIT_RATE)) {
+                            bitRate = decoderFormat.getInteger(MediaFormat.KEY_BIT_RATE);
+                        } else {
+                            bitRate = ENCODE_DEFAULT_AUDIO_BIT_RATE;
+                        }
                     }
                     /*Setup Encode Format*/
                     MediaFormat encodeFormat;
@@ -219,6 +233,7 @@
                         encodeFormat.setInteger(MediaFormat.KEY_LEVEL, level);
                         encodeFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);
                         encodeFormat.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, frameSize);
+                        encodeFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, colorFormat);
                     } else {
                         encodeFormat = MediaFormat.createAudioFormat(mime, sampleRate, numChannels);
                         encodeFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitRate);
@@ -255,7 +270,7 @@
         fileInput.close();
     }
 
-    @Test
+    @Test(timeout = PER_TEST_TIMEOUT_MS)
     public void testNativeEncoder() throws Exception {
         File inputFile = new File(mInputFilePath + mInputFile);
         assertTrue("Cannot find " + mInputFile + " in directory " + mInputFilePath,
@@ -275,8 +290,8 @@
             // Encoding the decoder's output
             for (String codecName : mediaCodecs) {
                 Native nativeEncoder = new Native();
-                int status = nativeEncoder.Encode(
-                        mInputFilePath, mInputFile, mDecodedFile, mStatsFile, codecName);
+                int status = nativeEncoder
+                        .Encode(mInputFilePath, mInputFile, mDecodedFile, mStatsFile, codecName);
                 assertEquals(
                         codecName + " encoder returned error " + status + " for " + "file:" + " " +
                                 mInputFile, 0, status);
diff --git a/media/tests/benchmark/MediaBenchmarkTest/src/main/cpp/NativeEncoder.cpp b/media/tests/benchmark/MediaBenchmarkTest/src/main/cpp/NativeEncoder.cpp
index 271b852..1277c8b 100644
--- a/media/tests/benchmark/MediaBenchmarkTest/src/main/cpp/NativeEncoder.cpp
+++ b/media/tests/benchmark/MediaBenchmarkTest/src/main/cpp/NativeEncoder.cpp
@@ -18,9 +18,9 @@
 #define LOG_TAG "NativeEncoder"
 
 #include <jni.h>
+#include <sys/stat.h>
 #include <fstream>
 #include <iostream>
-#include <sys/stat.h>
 
 #include <android/log.h>
 
@@ -29,6 +29,11 @@
 
 #include <stdio.h>
 
+constexpr int32_t ENCODE_DEFAULT_FRAME_RATE = 25;
+constexpr int32_t ENCODE_DEFAULT_AUDIO_BIT_RATE = 128000 /* 128 Kbps */;
+constexpr int32_t ENCODE_DEFAULT_BIT_RATE = 8000000 /* 8 Mbps */;
+constexpr int32_t ENCODE_MIN_BIT_RATE = 600000 /* 600 Kbps */;
+
 extern "C" JNIEXPORT int JNICALL Java_com_android_media_benchmark_library_Native_Encode(
         JNIEnv *env, jobject thiz, jstring jFilePath, jstring jFileName, jstring jOutFilePath,
         jstring jStatsFile, jstring jCodecName) {
@@ -72,7 +77,7 @@
             ALOGE("Track Format invalid");
             return -1;
         }
-        uint8_t *inputBuffer = (uint8_t *) malloc(fileSize);
+        uint8_t *inputBuffer = (uint8_t *)malloc(fileSize);
         if (!inputBuffer) {
             ALOGE("Insufficient memory");
             return -1;
@@ -110,6 +115,8 @@
             free(inputBuffer);
             return -1;
         }
+
+        AMediaFormat *decoderFormat = decoder->getFormat();
         AMediaFormat *format = extractor->getFormat();
         if (inputBuffer) {
             free(inputBuffer);
@@ -146,29 +153,34 @@
                 AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_FRAME_RATE, &encParams.frameRate);
                 AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_BIT_RATE, &encParams.bitrate);
                 if (encParams.bitrate <= 0 || encParams.frameRate <= 0) {
-                    encParams.frameRate = 25;
+                    encParams.frameRate = ENCODE_DEFAULT_FRAME_RATE;
                     if (!strcmp(mime, "video/3gpp") || !strcmp(mime, "video/mp4v-es")) {
-                        encParams.bitrate = 600000 /* 600 Kbps */;
+                        encParams.bitrate = ENCODE_MIN_BIT_RATE /* 600 Kbps */;
                     } else {
-                        encParams.bitrate = 8000000 /* 8 Mbps */;
+                        encParams.bitrate = ENCODE_DEFAULT_BIT_RATE /* 8 Mbps */;
                     }
                 }
                 AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_PROFILE, &encParams.profile);
                 AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_LEVEL, &encParams.level);
+                AMediaFormat_getInt32(decoderFormat, AMEDIAFORMAT_KEY_COLOR_FORMAT,
+                                      &encParams.colorFormat);
             } else {
                 AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_SAMPLE_RATE, &encParams.sampleRate);
                 AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_CHANNEL_COUNT,
                                       &encParams.numChannels);
-                encParams.bitrate =
-                        encParams.sampleRate * encParams.numChannels * 16 /* bitsPerSample */;
+                encParams.bitrate = ENCODE_DEFAULT_AUDIO_BIT_RATE;
             }
             Encoder *encoder = new Encoder();
             encoder->setupEncoder();
             status = encoder->encode(sCodecName, eleStream, eleSize, asyncMode[i], encParams,
-                                     (char *) mime);
+                                     (char *)mime);
+            if (status != AMEDIA_OK) {
+                ALOGE("Encoder returned error");
+                return -1;
+            }
+            ALOGV("Encoding complete with codec %s for asyncMode = %d", sCodecName.c_str(),
+                  asyncMode[i]);
             encoder->deInitCodec();
-            cout << "codec : " << codecName << endl;
-            ALOGV(" asyncMode = %d \n", asyncMode[i]);
             const char *statsFile = env->GetStringUTFChars(jStatsFile, nullptr);
             encoder->dumpStatistics(sInputReference, extractor->getClipDuration(), sCodecName,
                                     (asyncMode[i] ? "async" : "sync"), statsFile);
@@ -189,6 +201,10 @@
             AMediaFormat_delete(format);
             format = nullptr;
         }
+        if (decoderFormat) {
+            AMediaFormat_delete(decoderFormat);
+            decoderFormat = nullptr;
+        }
         decoder->deInitCodec();
         decoder->resetDecoder();
     }
diff --git a/media/tests/benchmark/MediaBenchmarkTest/src/main/java/com/android/media/benchmark/library/Decoder.java b/media/tests/benchmark/MediaBenchmarkTest/src/main/java/com/android/media/benchmark/library/Decoder.java
index 3b1eed4..c138350 100644
--- a/media/tests/benchmark/MediaBenchmarkTest/src/main/java/com/android/media/benchmark/library/Decoder.java
+++ b/media/tests/benchmark/MediaBenchmarkTest/src/main/java/com/android/media/benchmark/library/Decoder.java
@@ -256,6 +256,13 @@
      */
     public void resetDecoder() { mStats.reset(); }
 
+    /**
+     * Returns the format of the output buffers
+     */
+    public MediaFormat getFormat() {
+        return mCodec.getOutputFormat();
+    }
+
     private void onInputAvailable(int inputBufferId, MediaCodec mediaCodec) {
         if ((inputBufferId >= 0) && !mSawInputEOS) {
             ByteBuffer inputCodecBuffer = mediaCodec.getInputBuffer(inputBufferId);
diff --git a/media/tests/benchmark/src/native/decoder/Decoder.cpp b/media/tests/benchmark/src/native/decoder/Decoder.cpp
index 2171589..090f3e1 100644
--- a/media/tests/benchmark/src/native/decoder/Decoder.cpp
+++ b/media/tests/benchmark/src/native/decoder/Decoder.cpp
@@ -63,8 +63,8 @@
         ssize_t bytesRead = 0;
         uint32_t flag = 0;
         int64_t presentationTimeUs = 0;
-        tie(bytesRead, flag, presentationTimeUs) = readSampleData(
-                mInputBuffer, mOffset, mFrameMetaData, buf, mNumInputFrame, bufSize);
+        tie(bytesRead, flag, presentationTimeUs) =
+                readSampleData(mInputBuffer, mOffset, mFrameMetaData, buf, mNumInputFrame, bufSize);
         if (flag == AMEDIA_ERROR_MALFORMED) {
             mErrorCode = (media_status_t)flag;
             mSignalledError = true;
@@ -144,6 +144,11 @@
     if (!mFormat) mFormat = mExtractor->getFormat();
 }
 
+AMediaFormat *Decoder::getFormat() {
+    ALOGV("In %s", __func__);
+    return AMediaCodec_getOutputFormat(mCodec);
+}
+
 int32_t Decoder::decode(uint8_t *inputBuffer, vector<AMediaCodecBufferInfo> &frameInfo,
                         string &codecName, bool asyncMode, FILE *outFp) {
     ALOGV("In %s", __func__);
diff --git a/media/tests/benchmark/src/native/decoder/Decoder.h b/media/tests/benchmark/src/native/decoder/Decoder.h
index f3fa6a1..e619cb4 100644
--- a/media/tests/benchmark/src/native/decoder/Decoder.h
+++ b/media/tests/benchmark/src/native/decoder/Decoder.h
@@ -57,6 +57,8 @@
 
     void resetDecoder();
 
+    AMediaFormat *getFormat();
+
     // Async callback APIs
     void onInputAvailable(AMediaCodec *codec, int32_t index) override;
 
diff --git a/media/tests/benchmark/src/native/encoder/Encoder.cpp b/media/tests/benchmark/src/native/encoder/Encoder.cpp
index 2db612c..8dfe993 100644
--- a/media/tests/benchmark/src/native/encoder/Encoder.cpp
+++ b/media/tests/benchmark/src/native/encoder/Encoder.cpp
@@ -181,8 +181,8 @@
     mStats->dumpStatistics(operation, inputReference, durationUs, componentName, mode, statsFile);
 }
 
-int32_t Encoder::encode(string &codecName, ifstream &eleStream, size_t eleSize,
-                        bool asyncMode, encParameter encParams, char *mime) {
+int32_t Encoder::encode(string &codecName, ifstream &eleStream, size_t eleSize, bool asyncMode,
+                        encParameter encParams, char *mime) {
     ALOGV("In %s", __func__);
     mEleStream = &eleStream;
     mInputBufferSize = eleSize;
@@ -202,6 +202,7 @@
             AMediaFormat_setInt32(mFormat, AMEDIAFORMAT_KEY_PROFILE, mParams.profile);
             AMediaFormat_setInt32(mFormat, AMEDIAFORMAT_KEY_LEVEL, mParams.level);
         }
+        AMediaFormat_setInt32(mFormat, AMEDIAFORMAT_KEY_COLOR_FORMAT, mParams.colorFormat);
     } else {
         AMediaFormat_setInt32(mFormat, AMEDIAFORMAT_KEY_SAMPLE_RATE, mParams.sampleRate);
         AMediaFormat_setInt32(mFormat, AMEDIAFORMAT_KEY_CHANNEL_COUNT, mParams.numChannels);
diff --git a/media/tests/benchmark/src/native/encoder/Encoder.h b/media/tests/benchmark/src/native/encoder/Encoder.h
index 3d12600..5ad142b 100644
--- a/media/tests/benchmark/src/native/encoder/Encoder.h
+++ b/media/tests/benchmark/src/native/encoder/Encoder.h
@@ -23,9 +23,11 @@
 #include <queue>
 #include <thread>
 
+#include "media/NdkImage.h"
 #include "BenchmarkCommon.h"
 #include "Stats.h"
 
+
 struct encParameter {
     int32_t bitrate = -1;
     int32_t numFrames = -1;
@@ -38,6 +40,7 @@
     int32_t frameRate = -1;
     int32_t profile = 0;
     int32_t level = 0;
+    int32_t colorFormat = AIMAGE_FORMAT_YUV_420_888;
 };
 
 class Encoder : public CallBackHandle {