Use TrackSelector in DownloadHelper.

This adds the basic track selection capabilties (including tests).
The new capabilities are not exposed yet through the DownloadHelper implementations
and there will also be more helper methods (e.g. to select multiple audio lanuages at
once).

PiperOrigin-RevId: 224518477

library/core/src/main/java/com/google/android/exoplayer2/offline/DownloadHelper.java

@@ -19,18 +19,67 @@ import android.net.Uri;
 import android.os.Handler;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Looper;
 import android.support.annotation.Nullable;
 import android.support.annotation.Nullable;
+import android.util.SparseIntArray;
+import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.ExoPlaybackException;
+import com.google.android.exoplayer2.Renderer;
+import com.google.android.exoplayer2.RendererCapabilities;
+import com.google.android.exoplayer2.RenderersFactory;
+import com.google.android.exoplayer2.Timeline;
+import com.google.android.exoplayer2.drm.DrmSessionManager;
+import com.google.android.exoplayer2.drm.FrameworkMediaCrypto;
+import com.google.android.exoplayer2.source.MediaSource.MediaPeriodId;
+import com.google.android.exoplayer2.source.TrackGroup;
 import com.google.android.exoplayer2.source.TrackGroupArray;
 import com.google.android.exoplayer2.source.TrackGroupArray;
+import com.google.android.exoplayer2.trackselection.BaseTrackSelection;
+import com.google.android.exoplayer2.trackselection.DefaultTrackSelector;
+import com.google.android.exoplayer2.trackselection.DefaultTrackSelector.Parameters;
+import com.google.android.exoplayer2.trackselection.MappingTrackSelector.MappedTrackInfo;
+import com.google.android.exoplayer2.trackselection.TrackSelection;
+import com.google.android.exoplayer2.trackselection.TrackSelectorResult;
+import com.google.android.exoplayer2.upstream.BandwidthMeter;
+import com.google.android.exoplayer2.upstream.TransferListener;
 import com.google.android.exoplayer2.util.Assertions;
 import com.google.android.exoplayer2.util.Assertions;
+import com.google.android.exoplayer2.util.Util;
 import java.io.IOException;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.List;
+import org.checkerframework.checker.nullness.compatqual.NullableType;
+import org.checkerframework.checker.nullness.qual.EnsuresNonNull;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.RequiresNonNull;
 
 
 /**
 /**
  * A helper for initializing and removing downloads.
  * A helper for initializing and removing downloads.
  *
  *
+ * <p>The helper extracts track information from the media, selects tracks for downloading, and
+ * creates {@link DownloadAction download actions} based on the selected tracks.
+ *
+ * <p>A typical usage of DownloadHelper follows these steps:
+ *
+ * <ol>
+ *   <li>Construct the download helper with information about the {@link RenderersFactory renderers}
+ *       and {@link DefaultTrackSelector.Parameters parameters} for track selection.
+ *   <li>Prepare the helper using {@link #prepare(Callback)} and wait for the callback.
+ *   <li>Optional: Inspect the selected tracks using {@link #getMappedTrackInfo(int)} and {@link
+ *       #getTrackSelections(int, int)}, and make adjustments using {@link
+ *       #clearTrackSelections(int)}, {@link #replaceTrackSelections(int, Parameters)} and {@link
+ *       #addTrackSelection(int, Parameters)}.
+ *   <li>Create download actions for the selected track using {@link #getDownloadAction(byte[])}.
+ * </ol>
+ *
  * @param <T> The manifest type.
  * @param <T> The manifest type.
  */
  */
 public abstract class DownloadHelper<T> {
 public abstract class DownloadHelper<T> {
 
 
+  /**
+   * The default parameters used for track selection for downloading. This default selects the
+   * highest bitrate audio and video tracks which are supported by the renderers.
+   */
+  public static final DefaultTrackSelector.Parameters DEFAULT_TRACK_SELECTOR_PARAMETERS =
+      new DefaultTrackSelector.ParametersBuilder().setForceHighestSupportedBitrate(true).build();
+
   /** A callback to be notified when the {@link DownloadHelper} is prepared. */
   /** A callback to be notified when the {@link DownloadHelper} is prepared. */
   public interface Callback {
   public interface Callback {
 
 
@@ -39,7 +88,7 @@ public abstract class DownloadHelper<T> {
      *
      *
      * @param helper The reporting {@link DownloadHelper}.
      * @param helper The reporting {@link DownloadHelper}.
      */
      */
-    void onPrepared(DownloadHelper helper);
+    void onPrepared(DownloadHelper<?> helper);
 
 
     /**
     /**
      * Called when preparation fails.
      * Called when preparation fails.
@@ -47,15 +96,22 @@ public abstract class DownloadHelper<T> {
      * @param helper The reporting {@link DownloadHelper}.
      * @param helper The reporting {@link DownloadHelper}.
      * @param e The error.
      * @param e The error.
      */
      */
-    void onPrepareError(DownloadHelper helper, IOException e);
+    void onPrepareError(DownloadHelper<?> helper, IOException e);
   }
   }
 
 
   private final String downloadType;
   private final String downloadType;
   private final Uri uri;
   private final Uri uri;
   @Nullable private final String cacheKey;
   @Nullable private final String cacheKey;
+  private final DefaultTrackSelector trackSelector;
+  private final RendererCapabilities[] rendererCapabilities;
+  private final SparseIntArray scratchSet;
 
 
+  private int currentTrackSelectionPeriodIndex;
   @Nullable private T manifest;
   @Nullable private T manifest;
-  @Nullable private TrackGroupArray[] trackGroupArrays;
+  private TrackGroupArray @MonotonicNonNull [] trackGroupArrays;
+  private MappedTrackInfo @MonotonicNonNull [] mappedTrackInfos;
+  private List<TrackSelection> @MonotonicNonNull [][] trackSelectionsByPeriodAndRenderer;
+  private List<TrackSelection> @MonotonicNonNull [][] immutableTrackSelectionsByPeriodAndRenderer;
 
 
   /**
   /**
    * Create download helper.
    * Create download helper.
@@ -65,9 +121,45 @@ public abstract class DownloadHelper<T> {
    * @param cacheKey An optional cache key.
    * @param cacheKey An optional cache key.
    */
    */
   public DownloadHelper(String downloadType, Uri uri, @Nullable String cacheKey) {
   public DownloadHelper(String downloadType, Uri uri, @Nullable String cacheKey) {
+    // TODO: Remove as soon as all implementations have been updated to the new constructor.
+    this(
+        downloadType,
+        uri,
+        cacheKey,
+        DEFAULT_TRACK_SELECTOR_PARAMETERS,
+        /* renderersFactory= */ (handler, videoListener, audioListener, metadata, text, drm) ->
+            new Renderer[0],
+        /* drmSessionManager= */ null);
+  }
+
+  /**
+   * Creates download helper.
+   *
+   * @param downloadType A download type. This value will be used as {@link DownloadAction#type}.
+   * @param uri A {@link Uri}.
+   * @param cacheKey An optional cache key.
+   * @param trackSelectorParameters {@link DefaultTrackSelector.Parameters} for selecting tracks for
+   *     downloading.
+   * @param renderersFactory The {@link RenderersFactory} creating the renderers for which tracks
+   *     are selected.
+   * @param drmSessionManager An optional {@link DrmSessionManager} used by the renderers created by
+   *     {@code renderersFactory}.
+   */
+  public DownloadHelper(
+      String downloadType,
+      Uri uri,
+      @Nullable String cacheKey,
+      DefaultTrackSelector.Parameters trackSelectorParameters,
+      RenderersFactory renderersFactory,
+      @Nullable DrmSessionManager<FrameworkMediaCrypto> drmSessionManager) {
     this.downloadType = downloadType;
     this.downloadType = downloadType;
     this.uri = uri;
     this.uri = uri;
     this.cacheKey = cacheKey;
     this.cacheKey = cacheKey;
+    this.trackSelector = new DefaultTrackSelector(new DownloadTrackSelection.Factory());
+    this.rendererCapabilities = Util.getRendererCapabilities(renderersFactory, drmSessionManager);
+    this.scratchSet = new SparseIntArray();
+    trackSelector.setParameters(trackSelectorParameters);
+    trackSelector.init(/* listener= */ () -> {}, new DummyBandwidthMeter());
   }
   }
 
 
   /**
   /**
@@ -77,21 +169,28 @@ public abstract class DownloadHelper<T> {
    *     will be invoked on the calling thread unless that thread does not have an associated {@link
    *     will be invoked on the calling thread unless that thread does not have an associated {@link
    *     Looper}, in which case it will be called on the application's main thread.
    *     Looper}, in which case it will be called on the application's main thread.
    */
    */
-  public final void prepare(final Callback callback) {
-    final Handler handler =
+  public final void prepare(Callback callback) {
+    Handler handler =
         new Handler(Looper.myLooper() != null ? Looper.myLooper() : Looper.getMainLooper());
         new Handler(Looper.myLooper() != null ? Looper.myLooper() : Looper.getMainLooper());
-    new Thread() {
-      @Override
-      public void run() {
-        try {
-          manifest = loadManifest(uri);
-          trackGroupArrays = getTrackGroupArrays(manifest);
-          handler.post(() -> callback.onPrepared(DownloadHelper.this));
-        } catch (final IOException e) {
-          handler.post(() -> callback.onPrepareError(DownloadHelper.this, e));
-        }
-      }
-    }.start();
+    new Thread(
+            () -> {
+              try {
+                manifest = loadManifest(uri);
+                trackGroupArrays = getTrackGroupArrays(manifest);
+                initializeTrackSelectionLists(trackGroupArrays.length, rendererCapabilities.length);
+                mappedTrackInfos = new MappedTrackInfo[trackGroupArrays.length];
+                for (int i = 0; i < trackGroupArrays.length; i++) {
+                  TrackSelectorResult trackSelectorResult = runTrackSelection(/* periodIndex= */ i);
+                  trackSelector.onSelectionActivated(trackSelectorResult.info);
+                  mappedTrackInfos[i] =
+                      Assertions.checkNotNull(trackSelector.getCurrentMappedTrackInfo());
+                }
+                handler.post(() -> callback.onPrepared(DownloadHelper.this));
+              } catch (final IOException e) {
+                handler.post(() -> callback.onPrepareError(DownloadHelper.this, e));
+              }
+            })
+        .start();
   }
   }
 
 
   /** Returns the manifest. Must not be called until after preparation completes. */
   /** Returns the manifest. Must not be called until after preparation completes. */
@@ -113,6 +212,8 @@ public abstract class DownloadHelper<T> {
    * Returns the track groups for the given period. Must not be called until after preparation
    * Returns the track groups for the given period. Must not be called until after preparation
    * completes.
    * completes.
    *
    *
+   * <p>Use {@link #getMappedTrackInfo(int)} to get the track groups mapped to renderers.
+   *
    * @param periodIndex The period index.
    * @param periodIndex The period index.
    * @return The track groups for the period. May be {@link TrackGroupArray#EMPTY} for single stream
    * @return The track groups for the period. May be {@link TrackGroupArray#EMPTY} for single stream
    *     content.
    *     content.
@@ -123,6 +224,107 @@ public abstract class DownloadHelper<T> {
   }
   }
 
 
   /**
   /**
+   * Returns the mapped track info for the given period. Must not be called until after preparation
+   * completes.
+   *
+   * @param periodIndex The period index.
+   * @return The {@link MappedTrackInfo} for the period.
+   */
+  public final MappedTrackInfo getMappedTrackInfo(int periodIndex) {
+    Assertions.checkNotNull(mappedTrackInfos);
+    return mappedTrackInfos[periodIndex];
+  }
+
+  /**
+   * Returns all {@link TrackSelection track selections} for a period and renderer. Must not be
+   * called until after preparation completes.
+   *
+   * @param periodIndex The period index.
+   * @param rendererIndex The renderer index.
+   * @return A list of selected {@link TrackSelection track selections}.
+   */
+  public final List<TrackSelection> getTrackSelections(int periodIndex, int rendererIndex) {
+    Assertions.checkNotNull(immutableTrackSelectionsByPeriodAndRenderer);
+    return immutableTrackSelectionsByPeriodAndRenderer[periodIndex][rendererIndex];
+  }
+
+  /**
+   * Clears the selection of tracks for a period. Must not be called until after preparation
+   * completes.
+   *
+   * @param periodIndex The period index for which track selections are cleared.
+   */
+  public final void clearTrackSelections(int periodIndex) {
+    Assertions.checkNotNull(trackSelectionsByPeriodAndRenderer);
+    for (int i = 0; i < rendererCapabilities.length; i++) {
+      trackSelectionsByPeriodAndRenderer[periodIndex][i].clear();
+    }
+  }
+
+  /**
+   * Replaces a selection of tracks to be downloaded. Must not be called until after preparation
+   * completes.
+   *
+   * @param periodIndex The period index for which the track selection is replaced.
+   * @param trackSelectorParameters The {@link DefaultTrackSelector.Parameters} to obtain the new
+   *     selection of tracks.
+   */
+  public final void replaceTrackSelections(
+      int periodIndex, DefaultTrackSelector.Parameters trackSelectorParameters) {
+    clearTrackSelections(periodIndex);
+    addTrackSelection(periodIndex, trackSelectorParameters);
+  }
+
+  /**
+   * Adds a selection of tracks to be downloaded. Must not be called until after preparation
+   * completes.
+   *
+   * @param periodIndex The period index this track selection is added for.
+   * @param trackSelectorParameters The {@link DefaultTrackSelector.Parameters} to obtain the new
+   *     selection of tracks.
+   */
+  public final void addTrackSelection(
+      int periodIndex, DefaultTrackSelector.Parameters trackSelectorParameters) {
+    Assertions.checkNotNull(trackGroupArrays);
+    Assertions.checkNotNull(trackSelectionsByPeriodAndRenderer);
+    trackSelector.setParameters(trackSelectorParameters);
+    runTrackSelection(periodIndex);
+  }
+
+  /**
+   * Builds a {@link DownloadAction} for downloading the selected tracks. Must not be called until
+   * after preparation completes.
+   *
+   * @param data Application provided data to store in {@link DownloadAction#data}.
+   * @return The built {@link DownloadAction}.
+   */
+  public final DownloadAction getDownloadAction(@Nullable byte[] data) {
+    Assertions.checkNotNull(trackSelectionsByPeriodAndRenderer);
+    Assertions.checkNotNull(trackGroupArrays);
+    List<TrackKey> trackKeys = new ArrayList<>();
+    int periodCount = trackSelectionsByPeriodAndRenderer.length;
+    for (int periodIndex = 0; periodIndex < periodCount; periodIndex++) {
+      int rendererCount = trackSelectionsByPeriodAndRenderer[periodIndex].length;
+      for (int rendererIndex = 0; rendererIndex < rendererCount; rendererIndex++) {
+        List<TrackSelection> trackSelectionList =
+            trackSelectionsByPeriodAndRenderer[periodIndex][rendererIndex];
+        for (int selectionIndex = 0; selectionIndex < trackSelectionList.size(); selectionIndex++) {
+          TrackSelection trackSelection = trackSelectionList.get(selectionIndex);
+          int trackGroupIndex =
+              trackGroupArrays[periodIndex].indexOf(trackSelection.getTrackGroup());
+          int trackCount = trackSelection.length();
+          for (int trackListIndex = 0; trackListIndex < trackCount; trackListIndex++) {
+            int trackIndex = trackSelection.getIndexInTrackGroup(trackListIndex);
+            trackKeys.add(new TrackKey(periodIndex, trackGroupIndex, trackIndex));
+          }
+        }
+      }
+    }
+    return DownloadAction.createDownloadAction(
+        downloadType, uri, toStreamKeys(trackKeys), cacheKey, data);
+  }
+
+  /**
    * Builds a {@link DownloadAction} for downloading the specified tracks. Must not be called until
    * Builds a {@link DownloadAction} for downloading the specified tracks. Must not be called until
    * after preparation completes.
    * after preparation completes.
    *
    *
@@ -131,6 +333,7 @@ public abstract class DownloadHelper<T> {
    * @return The built {@link DownloadAction}.
    * @return The built {@link DownloadAction}.
    */
    */
   public final DownloadAction getDownloadAction(@Nullable byte[] data, List<TrackKey> trackKeys) {
   public final DownloadAction getDownloadAction(@Nullable byte[] data, List<TrackKey> trackKeys) {
+    // TODO: Remove as soon as all usages have been updated to new getDownloadAction method.
     return DownloadAction.createDownloadAction(
     return DownloadAction.createDownloadAction(
         downloadType, uri, toStreamKeys(trackKeys), cacheKey, data);
         downloadType, uri, toStreamKeys(trackKeys), cacheKey, data);
   }
   }
@@ -167,4 +370,142 @@ public abstract class DownloadHelper<T> {
    * @return A corresponding list of stream keys.
    * @return A corresponding list of stream keys.
    */
    */
   protected abstract List<StreamKey> toStreamKeys(List<TrackKey> trackKeys);
   protected abstract List<StreamKey> toStreamKeys(List<TrackKey> trackKeys);
+
+  @SuppressWarnings("unchecked")
+  @EnsuresNonNull("trackSelectionsByPeriodAndRenderer")
+  private void initializeTrackSelectionLists(int periodCount, int rendererCount) {
+    trackSelectionsByPeriodAndRenderer =
+        (List<TrackSelection>[][]) new List<?>[periodCount][rendererCount];
+    immutableTrackSelectionsByPeriodAndRenderer =
+        (List<TrackSelection>[][]) new List<?>[periodCount][rendererCount];
+    for (int i = 0; i < periodCount; i++) {
+      for (int j = 0; j < rendererCount; j++) {
+        trackSelectionsByPeriodAndRenderer[i][j] = new ArrayList<>();
+        immutableTrackSelectionsByPeriodAndRenderer[i][j] =
+            Collections.unmodifiableList(trackSelectionsByPeriodAndRenderer[i][j]);
+      }
+    }
+  }
+
+  /**
+   * Runs the track selection for a given period index with the current parameters. The selected
+   * tracks will be added to {@link #trackSelectionsByPeriodAndRenderer}.
+   */
+  // Intentional reference comparison of track group instances.
+  @SuppressWarnings("ReferenceEquality")
+  @RequiresNonNull({"trackGroupArrays", "trackSelectionsByPeriodAndRenderer"})
+  private TrackSelectorResult runTrackSelection(int periodIndex) {
+    // TODO: Use actual timeline and media period id.
+    MediaPeriodId dummyMediaPeriodId = new MediaPeriodId(new Object());
+    Timeline dummyTimeline = Timeline.EMPTY;
+    currentTrackSelectionPeriodIndex = periodIndex;
+    try {
+      TrackSelectorResult trackSelectorResult =
+          trackSelector.selectTracks(
+              rendererCapabilities,
+              trackGroupArrays[periodIndex],
+              dummyMediaPeriodId,
+              dummyTimeline);
+      for (int i = 0; i < trackSelectorResult.length; i++) {
+        TrackSelection newSelection = trackSelectorResult.selections.get(i);
+        if (newSelection == null) {
+          continue;
+        }
+        List<TrackSelection> existingSelectionList =
+            trackSelectionsByPeriodAndRenderer[currentTrackSelectionPeriodIndex][i];
+        boolean mergedWithExistingSelection = false;
+        for (int j = 0; j < existingSelectionList.size(); j++) {
+          TrackSelection existingSelection = existingSelectionList.get(j);
+          if (existingSelection.getTrackGroup() == newSelection.getTrackGroup()) {
+            // Merge with existing selection.
+            scratchSet.clear();
+            for (int k = 0; k < existingSelection.length(); k++) {
+              scratchSet.put(existingSelection.getIndexInTrackGroup(k), 0);
+            }
+            for (int k = 0; k < newSelection.length(); k++) {
+              scratchSet.put(newSelection.getIndexInTrackGroup(k), 0);
+            }
+            int[] mergedTracks = new int[scratchSet.size()];
+            for (int k = 0; k < scratchSet.size(); k++) {
+              mergedTracks[k] = scratchSet.keyAt(k);
+            }
+            existingSelectionList.set(
+                j, new DownloadTrackSelection(existingSelection.getTrackGroup(), mergedTracks));
+            mergedWithExistingSelection = true;
+            break;
+          }
+        }
+        if (!mergedWithExistingSelection) {
+          existingSelectionList.add(newSelection);
+        }
+      }
+      return trackSelectorResult;
+    } catch (ExoPlaybackException e) {
+      // DefaultTrackSelector does not throw exceptions during track selection.
+      throw new UnsupportedOperationException(e);
+    }
+  }
+
+  private static final class DownloadTrackSelection extends BaseTrackSelection {
+
+    private static final class Factory implements TrackSelection.Factory {
+
+      @Override
+      public @NullableType TrackSelection[] createTrackSelections(
+          @NullableType Definition[] definitions, BandwidthMeter bandwidthMeter) {
+        @NullableType TrackSelection[] selections = new TrackSelection[definitions.length];
+        for (int i = 0; i < definitions.length; i++) {
+          selections[i] =
+              definitions[i] == null
+                  ? null
+                  : new DownloadTrackSelection(definitions[i].group, definitions[i].tracks);
+        }
+        return selections;
+      }
+    }
+
+    public DownloadTrackSelection(TrackGroup trackGroup, int[] tracks) {
+      super(trackGroup, tracks);
+    }
+
+    @Override
+    public int getSelectedIndex() {
+      return 0;
+    }
+
+    @Override
+    public int getSelectionReason() {
+      return C.SELECTION_REASON_UNKNOWN;
+    }
+
+    @Nullable
+    @Override
+    public Object getSelectionData() {
+      return null;
+    }
+  }
+
+  private static final class DummyBandwidthMeter implements BandwidthMeter {
+
+    @Override
+    public long getBitrateEstimate() {
+      return 0;
+    }
+
+    @Nullable
+    @Override
+    public TransferListener getTransferListener() {
+      return null;
+    }
+
+    @Override
+    public void addEventListener(Handler eventHandler, EventListener eventListener) {
+      // Do nothing.
+    }
+
+    @Override
+    public void removeEventListener(EventListener eventListener) {
+      // Do nothing.
+    }
+  }
 }
 }

library/core/src/main/java/com/google/android/exoplayer2/offline/TrackKey.java

@@ -15,6 +15,8 @@
  */
  */
 package com.google.android.exoplayer2.offline;
 package com.google.android.exoplayer2.offline;
 
 
+import android.support.annotation.Nullable;
+
 /**
 /**
  * Identifies a given track by the index of the containing period, the index of the containing group
  * Identifies a given track by the index of the containing period, the index of the containing group
  * within the period, and the index of the track within the group.
  * within the period, and the index of the track within the group.
@@ -38,4 +40,23 @@ public final class TrackKey {
     this.groupIndex = groupIndex;
     this.groupIndex = groupIndex;
     this.trackIndex = trackIndex;
     this.trackIndex = trackIndex;
   }
   }
+
+  @Override
+  public boolean equals(@Nullable Object other) {
+    if (this == other) {
+      return true;
+    }
+    if (other == null || getClass() != other.getClass()) {
+      return false;
+    }
+    TrackKey that = (TrackKey) other;
+    return periodIndex == that.periodIndex
+        && groupIndex == that.groupIndex
+        && trackIndex == that.trackIndex;
+  }
+
+  @Override
+  public int hashCode() {
+    return 31 * (31 * periodIndex + groupIndex) + trackIndex;
+  }
 }
 }

library/core/src/test/java/com/google/android/exoplayer2/offline/DownloadHelperTest.java

+/*
+ * Copyright (C) 2018 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 com.google.android.exoplayer2.offline;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.net.Uri;
+import android.support.annotation.Nullable;
+import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.Format;
+import com.google.android.exoplayer2.Renderer;
+import com.google.android.exoplayer2.RenderersFactory;
+import com.google.android.exoplayer2.offline.DownloadHelper.Callback;
+import com.google.android.exoplayer2.source.TrackGroup;
+import com.google.android.exoplayer2.source.TrackGroupArray;
+import com.google.android.exoplayer2.testutil.FakeRenderer;
+import com.google.android.exoplayer2.trackselection.DefaultTrackSelector;
+import com.google.android.exoplayer2.trackselection.DefaultTrackSelector.ParametersBuilder;
+import com.google.android.exoplayer2.trackselection.MappingTrackSelector.MappedTrackInfo;
+import com.google.android.exoplayer2.trackselection.TrackSelection;
+import com.google.android.exoplayer2.util.ConditionVariable;
+import com.google.android.exoplayer2.util.MimeTypes;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.shadows.ShadowLooper;
+
+/** Unit tests for {@link DownloadHelper}. */
+@RunWith(RobolectricTestRunner.class)
+public class DownloadHelperTest {
+
+  private static final String TEST_DOWNLOAD_TYPE = "downloadType";
+  private static final String TEST_CACHE_KEY = "cacheKey";
+  private static final ManifestType TEST_MANIFEST = new ManifestType();
+  private static final List<StreamKey> testStreamKeys =
+      Arrays.asList(new StreamKey(0, 1, 2), new StreamKey(1, 3, 4));
+
+  private static final Format VIDEO_FORMAT_LOW = createVideoFormat(/* bitrate= */ 200_000);
+  private static final Format VIDEO_FORMAT_HIGH = createVideoFormat(/* bitrate= */ 800_000);
+  private static final Format AUDIO_FORMAT_US = createAudioFormat(/* language= */ "US");
+  private static final Format AUDIO_FORMAT_ZH = createAudioFormat(/* language= */ "ZH");
+  private static final Format TEXT_FORMAT_US = createTextFormat(/* language= */ "US");
+  private static final Format TEXT_FORMAT_ZH = createTextFormat(/* language= */ "ZH");
+
+  private static final TrackGroup TRACK_GROUP_VIDEO_BOTH =
+      new TrackGroup(VIDEO_FORMAT_LOW, VIDEO_FORMAT_HIGH);
+  private static final TrackGroup TRACK_GROUP_VIDEO_SINGLE = new TrackGroup(VIDEO_FORMAT_LOW);
+  private static final TrackGroup TRACK_GROUP_AUDIO_US = new TrackGroup(AUDIO_FORMAT_US);
+  private static final TrackGroup TRACK_GROUP_AUDIO_ZH = new TrackGroup(AUDIO_FORMAT_ZH);
+  private static final TrackGroup TRACK_GROUP_TEXT_US = new TrackGroup(TEXT_FORMAT_US);
+  private static final TrackGroup TRACK_GROUP_TEXT_ZH = new TrackGroup(TEXT_FORMAT_ZH);
+  private static final TrackGroupArray TRACK_GROUP_ARRAY_ALL =
+      new TrackGroupArray(
+          TRACK_GROUP_VIDEO_BOTH,
+          TRACK_GROUP_AUDIO_US,
+          TRACK_GROUP_AUDIO_ZH,
+          TRACK_GROUP_TEXT_US,
+          TRACK_GROUP_TEXT_ZH);
+  private static final TrackGroupArray TRACK_GROUP_ARRAY_SINGLE =
+      new TrackGroupArray(TRACK_GROUP_VIDEO_SINGLE, TRACK_GROUP_AUDIO_US);
+  private static final TrackGroupArray[] TRACK_GROUP_ARRAYS =
+      new TrackGroupArray[] {TRACK_GROUP_ARRAY_ALL, TRACK_GROUP_ARRAY_SINGLE};
+
+  private Uri testUri;
+
+  private FakeDownloadHelper downloadHelper;
+
+  @Before
+  public void setUp() {
+    testUri = Uri.parse("http://test.uri");
+
+    FakeRenderer videoRenderer = new FakeRenderer(VIDEO_FORMAT_LOW, VIDEO_FORMAT_HIGH);
+    FakeRenderer audioRenderer = new FakeRenderer(AUDIO_FORMAT_US, AUDIO_FORMAT_ZH);
+    FakeRenderer textRenderer = new FakeRenderer(TEXT_FORMAT_US, TEXT_FORMAT_ZH);
+    RenderersFactory renderersFactory =
+        (handler, videoListener, audioListener, metadata, text, drm) ->
+            new Renderer[] {textRenderer, audioRenderer, videoRenderer};
+
+    downloadHelper = new FakeDownloadHelper(testUri, renderersFactory);
+  }
+
+  @Test
+  public void getManifest_returnsManifest() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+
+    ManifestType manifest = downloadHelper.getManifest();
+
+    assertThat(manifest).isEqualTo(TEST_MANIFEST);
+  }
+
+  @Test
+  public void getPeriodCount_returnsPeriodCount() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+
+    int periodCount = downloadHelper.getPeriodCount();
+
+    assertThat(periodCount).isEqualTo(2);
+  }
+
+  @Test
+  public void getTrackGroups_returnsTrackGroups() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+
+    TrackGroupArray trackGroupArrayPeriod0 = downloadHelper.getTrackGroups(/* periodIndex= */ 0);
+    TrackGroupArray trackGroupArrayPeriod1 = downloadHelper.getTrackGroups(/* periodIndex= */ 1);
+
+    assertThat(trackGroupArrayPeriod0).isEqualTo(TRACK_GROUP_ARRAYS[0]);
+    assertThat(trackGroupArrayPeriod1).isEqualTo(TRACK_GROUP_ARRAYS[1]);
+  }
+
+  @Test
+  public void getMappedTrackInfo_returnsMappedTrackInfo() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+
+    MappedTrackInfo mappedTracks0 = downloadHelper.getMappedTrackInfo(/* periodIndex= */ 0);
+    MappedTrackInfo mappedTracks1 = downloadHelper.getMappedTrackInfo(/* periodIndex= */ 1);
+
+    assertThat(mappedTracks0.getRendererCount()).isEqualTo(3);
+    assertThat(mappedTracks0.getRendererType(/* rendererIndex= */ 0)).isEqualTo(C.TRACK_TYPE_TEXT);
+    assertThat(mappedTracks0.getRendererType(/* rendererIndex= */ 1)).isEqualTo(C.TRACK_TYPE_AUDIO);
+    assertThat(mappedTracks0.getRendererType(/* rendererIndex= */ 2)).isEqualTo(C.TRACK_TYPE_VIDEO);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 0).length).isEqualTo(2);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 1).length).isEqualTo(2);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 2).length).isEqualTo(1);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 0).get(/* index= */ 0))
+        .isEqualTo(TRACK_GROUP_TEXT_US);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 0).get(/* index= */ 1))
+        .isEqualTo(TRACK_GROUP_TEXT_ZH);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 1).get(/* index= */ 0))
+        .isEqualTo(TRACK_GROUP_AUDIO_US);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 1).get(/* index= */ 1))
+        .isEqualTo(TRACK_GROUP_AUDIO_ZH);
+    assertThat(mappedTracks0.getTrackGroups(/* rendererIndex= */ 2).get(/* index= */ 0))
+        .isEqualTo(TRACK_GROUP_VIDEO_BOTH);
+
+    assertThat(mappedTracks1.getRendererCount()).isEqualTo(3);
+    assertThat(mappedTracks1.getRendererType(/* rendererIndex= */ 0)).isEqualTo(C.TRACK_TYPE_TEXT);
+    assertThat(mappedTracks1.getRendererType(/* rendererIndex= */ 1)).isEqualTo(C.TRACK_TYPE_AUDIO);
+    assertThat(mappedTracks1.getRendererType(/* rendererIndex= */ 2)).isEqualTo(C.TRACK_TYPE_VIDEO);
+    assertThat(mappedTracks1.getTrackGroups(/* rendererIndex= */ 0).length).isEqualTo(0);
+    assertThat(mappedTracks1.getTrackGroups(/* rendererIndex= */ 1).length).isEqualTo(1);
+    assertThat(mappedTracks1.getTrackGroups(/* rendererIndex= */ 2).length).isEqualTo(1);
+    assertThat(mappedTracks1.getTrackGroups(/* rendererIndex= */ 1).get(/* index= */ 0))
+        .isEqualTo(TRACK_GROUP_AUDIO_US);
+    assertThat(mappedTracks1.getTrackGroups(/* rendererIndex= */ 2).get(/* index= */ 0))
+        .isEqualTo(TRACK_GROUP_VIDEO_SINGLE);
+  }
+
+  @Test
+  public void getTrackSelections_returnsInitialSelection() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+
+    List<TrackSelection> selectedText0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 2);
+    List<TrackSelection> selectedText1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 2);
+
+    assertSingleTrackSelectionEquals(selectedText0, TRACK_GROUP_TEXT_US, 0);
+    assertSingleTrackSelectionEquals(selectedAudio0, TRACK_GROUP_AUDIO_US, 0);
+    assertSingleTrackSelectionEquals(selectedVideo0, TRACK_GROUP_VIDEO_BOTH, 1);
+
+    assertThat(selectedText1).isEmpty();
+    assertSingleTrackSelectionEquals(selectedAudio1, TRACK_GROUP_AUDIO_US, 0);
+    assertSingleTrackSelectionEquals(selectedVideo1, TRACK_GROUP_VIDEO_SINGLE, 0);
+  }
+
+  @Test
+  public void getTrackSelections_afterClearTrackSelections_isEmpty() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+
+    // Clear only one period selection to verify second period selection is untouched.
+    downloadHelper.clearTrackSelections(/* periodIndex= */ 0);
+    List<TrackSelection> selectedText0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 2);
+    List<TrackSelection> selectedText1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 2);
+
+    assertThat(selectedText0).isEmpty();
+    assertThat(selectedAudio0).isEmpty();
+    assertThat(selectedVideo0).isEmpty();
+
+    // Verify
+    assertThat(selectedText1).isEmpty();
+    assertSingleTrackSelectionEquals(selectedAudio1, TRACK_GROUP_AUDIO_US, 0);
+    assertSingleTrackSelectionEquals(selectedVideo1, TRACK_GROUP_VIDEO_SINGLE, 0);
+  }
+
+  @Test
+  public void getTrackSelections_afterReplaceTrackSelections_returnsNewSelections()
+      throws Exception {
+    prepareDownloadHelper(downloadHelper);
+    DefaultTrackSelector.Parameters parameters =
+        new ParametersBuilder()
+            .setPreferredAudioLanguage("ZH")
+            .setPreferredTextLanguage("ZH")
+            .setRendererDisabled(/* rendererIndex= */ 2, true)
+            .build();
+
+    // Replace only one period selection to verify second period selection is untouched.
+    downloadHelper.replaceTrackSelections(/* periodIndex= */ 0, parameters);
+    List<TrackSelection> selectedText0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 2);
+    List<TrackSelection> selectedText1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 2);
+
+    assertSingleTrackSelectionEquals(selectedText0, TRACK_GROUP_TEXT_ZH, 0);
+    assertSingleTrackSelectionEquals(selectedAudio0, TRACK_GROUP_AUDIO_ZH, 0);
+    assertThat(selectedVideo0).isEmpty();
+
+    assertThat(selectedText1).isEmpty();
+    assertSingleTrackSelectionEquals(selectedAudio1, TRACK_GROUP_AUDIO_US, 0);
+    assertSingleTrackSelectionEquals(selectedVideo1, TRACK_GROUP_VIDEO_SINGLE, 0);
+  }
+
+  @Test
+  public void getTrackSelections_afterAddTrackSelections_returnsCombinedSelections()
+      throws Exception {
+    prepareDownloadHelper(downloadHelper);
+    // Select parameters to require some merging of track groups because the new parameters add
+    // all video tracks to initial video single track selection.
+    DefaultTrackSelector.Parameters parameters =
+        new ParametersBuilder()
+            .setPreferredAudioLanguage("ZH")
+            .setPreferredTextLanguage("US")
+            .build();
+
+    // Add only to one period selection to verify second period selection is untouched.
+    downloadHelper.addTrackSelection(/* periodIndex= */ 0, parameters);
+    List<TrackSelection> selectedText0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo0 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 0, /* rendererIndex= */ 2);
+    List<TrackSelection> selectedText1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 0);
+    List<TrackSelection> selectedAudio1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 1);
+    List<TrackSelection> selectedVideo1 =
+        downloadHelper.getTrackSelections(/* periodIndex= */ 1, /* rendererIndex= */ 2);
+
+    assertSingleTrackSelectionEquals(selectedText0, TRACK_GROUP_TEXT_US, 0);
+    assertThat(selectedAudio0).hasSize(2);
+    assertTrackSelectionEquals(selectedAudio0.get(0), TRACK_GROUP_AUDIO_US, 0);
+    assertTrackSelectionEquals(selectedAudio0.get(1), TRACK_GROUP_AUDIO_ZH, 0);
+    assertSingleTrackSelectionEquals(selectedVideo0, TRACK_GROUP_VIDEO_BOTH, 0, 1);
+
+    assertThat(selectedText1).isEmpty();
+    assertSingleTrackSelectionEquals(selectedAudio1, TRACK_GROUP_AUDIO_US, 0);
+    assertSingleTrackSelectionEquals(selectedVideo1, TRACK_GROUP_VIDEO_SINGLE, 0);
+  }
+
+  @Test
+  public void getDownloadAction_createsDownloadAction_withAllSelectedTracks() throws Exception {
+    prepareDownloadHelper(downloadHelper);
+    // Ensure we have track groups with multiple indices, renderers with multiple track groups and
+    // also renderers without any track groups.
+    DefaultTrackSelector.Parameters parameters =
+        new ParametersBuilder()
+            .setPreferredAudioLanguage("ZH")
+            .setPreferredTextLanguage("US")
+            .build();
+    downloadHelper.addTrackSelection(/* periodIndex= */ 0, parameters);
+    byte[] data = new byte[10];
+    Arrays.fill(data, (byte) 123);
+
+    DownloadAction downloadAction = downloadHelper.getDownloadAction(data);
+
+    assertThat(downloadAction.type).isEqualTo(TEST_DOWNLOAD_TYPE);
+    assertThat(downloadAction.uri).isEqualTo(testUri);
+    assertThat(downloadAction.customCacheKey).isEqualTo(TEST_CACHE_KEY);
+    assertThat(downloadAction.isRemoveAction).isFalse();
+    assertThat(downloadAction.data).isEqualTo(data);
+    assertThat(downloadAction.keys).isEqualTo(testStreamKeys);
+    assertThat(downloadHelper.lastCreatedTrackKeys)
+        .containsExactly(
+            new TrackKey(/* periodIndex= */ 0, /* groupIndex= */ 0, /* trackIndex= */ 0),
+            new TrackKey(/* periodIndex= */ 0, /* groupIndex= */ 0, /* trackIndex= */ 1),
+            new TrackKey(/* periodIndex= */ 0, /* groupIndex= */ 1, /* trackIndex= */ 0),
+            new TrackKey(/* periodIndex= */ 0, /* groupIndex= */ 2, /* trackIndex= */ 0),
+            new TrackKey(/* periodIndex= */ 0, /* groupIndex= */ 3, /* trackIndex= */ 0),
+            new TrackKey(/* periodIndex= */ 1, /* groupIndex= */ 0, /* trackIndex= */ 0),
+            new TrackKey(/* periodIndex= */ 1, /* groupIndex= */ 1, /* trackIndex= */ 0));
+  }
+
+  @Test
+  public void getRemoveAction_returnsRemoveAction() {
+    DownloadAction removeAction = downloadHelper.getRemoveAction();
+
+    assertThat(removeAction.type).isEqualTo(TEST_DOWNLOAD_TYPE);
+    assertThat(removeAction.uri).isEqualTo(testUri);
+    assertThat(removeAction.customCacheKey).isEqualTo(TEST_CACHE_KEY);
+    assertThat(removeAction.isRemoveAction).isTrue();
+  }
+
+  private static void prepareDownloadHelper(FakeDownloadHelper downloadHelper) throws Exception {
+    AtomicReference<Exception> prepareException = new AtomicReference<>(null);
+    ConditionVariable preparedCondition = new ConditionVariable();
+    downloadHelper.prepare(
+        new Callback() {
+          @Override
+          public void onPrepared(DownloadHelper<?> helper) {
+            preparedCondition.open();
+          }
+
+          @Override
+          public void onPrepareError(DownloadHelper<?> helper, IOException e) {
+            prepareException.set(e);
+            preparedCondition.open();
+          }
+        });
+    while (!preparedCondition.block(0)) {
+      ShadowLooper.runMainLooperToNextTask();
+    }
+    if (prepareException.get() != null) {
+      throw prepareException.get();
+    }
+  }
+
+  private static Format createVideoFormat(int bitrate) {
+    return Format.createVideoSampleFormat(
+        /* id= */ null,
+        /* sampleMimeType= */ MimeTypes.VIDEO_H264,
+        /* codecs= */ null,
+        /* bitrate= */ bitrate,
+        /* maxInputSize= */ Format.NO_VALUE,
+        /* width= */ 480,
+        /* height= */ 360,
+        /* frameRate= */ Format.NO_VALUE,
+        /* initializationData= */ null,
+        /* drmInitData= */ null);
+  }
+
+  private static Format createAudioFormat(String language) {
+    return Format.createAudioSampleFormat(
+        /* id= */ null,
+        /* sampleMimeType= */ MimeTypes.AUDIO_AAC,
+        /* codecs= */ null,
+        /* bitrate= */ 48000,
+        /* maxInputSize= */ Format.NO_VALUE,
+        /* channelCount= */ 2,
+        /* sampleRate */ 44100,
+        /* initializationData= */ null,
+        /* drmInitData= */ null,
+        /* selectionFlags= */ C.SELECTION_FLAG_DEFAULT,
+        /* language= */ language);
+  }
+
+  private static Format createTextFormat(String language) {
+    return Format.createTextSampleFormat(
+        /* id= */ null,
+        /* sampleMimeType= */ MimeTypes.TEXT_VTT,
+        /* selectionFlags= */ C.SELECTION_FLAG_DEFAULT,
+        /* language= */ language);
+  }
+
+  private static void assertSingleTrackSelectionEquals(
+      List<TrackSelection> trackSelectionList, TrackGroup trackGroup, int... tracks) {
+    assertThat(trackSelectionList).hasSize(1);
+    assertTrackSelectionEquals(trackSelectionList.get(0), trackGroup, tracks);
+  }
+
+  private static void assertTrackSelectionEquals(
+      TrackSelection trackSelection, TrackGroup trackGroup, int... tracks) {
+    assertThat(trackSelection.getTrackGroup()).isEqualTo(trackGroup);
+    assertThat(trackSelection.length()).isEqualTo(tracks.length);
+    int[] selectedTracksInGroup = new int[trackSelection.length()];
+    for (int i = 0; i < trackSelection.length(); i++) {
+      selectedTracksInGroup[i] = trackSelection.getIndexInTrackGroup(i);
+    }
+    Arrays.sort(selectedTracksInGroup);
+    Arrays.sort(tracks);
+    assertThat(selectedTracksInGroup).isEqualTo(tracks);
+  }
+
+  private static final class ManifestType {}
+
+  private static final class FakeDownloadHelper extends DownloadHelper<ManifestType> {
+
+    @Nullable public List<TrackKey> lastCreatedTrackKeys;
+
+    public FakeDownloadHelper(Uri testUri, RenderersFactory renderersFactory) {
+      super(
+          TEST_DOWNLOAD_TYPE,
+          testUri,
+          TEST_CACHE_KEY,
+          DownloadHelper.DEFAULT_TRACK_SELECTOR_PARAMETERS,
+          renderersFactory,
+          /* drmSessionManager= */ null);
+    }
+
+    @Override
+    protected ManifestType loadManifest(Uri uri) throws IOException {
+      return TEST_MANIFEST;
+    }
+
+    @Override
+    protected TrackGroupArray[] getTrackGroupArrays(ManifestType manifest) {
+      assertThat(manifest).isEqualTo(TEST_MANIFEST);
+      return TRACK_GROUP_ARRAYS;
+    }
+
+    @Override
+    protected List<StreamKey> toStreamKeys(List<TrackKey> trackKeys) {
+      lastCreatedTrackKeys = trackKeys;
+      return testStreamKeys;
+    }
+  }
+}