/*
 * Copyright (C) 2020 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.exoplayer.drm;

import static androidx.media3.common.util.Assertions.checkNotNull;
import static com.google.common.truth.Truth.assertThat;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.junit.Assert.assertThrows;

import android.os.Looper;
import androidx.annotation.Nullable;
import androidx.media3.common.C;
import androidx.media3.common.DrmInitData;
import androidx.media3.common.Format;
import androidx.media3.common.MimeTypes;
import androidx.media3.common.util.Util;
import androidx.media3.exoplayer.analytics.PlayerId;
import androidx.media3.exoplayer.drm.DrmSessionManager.DrmSessionReference;
import androidx.media3.exoplayer.drm.ExoMediaDrm.AppManagedProvider;
import androidx.media3.exoplayer.source.MediaSource;
import androidx.media3.test.utils.FakeExoMediaDrm;
import androidx.media3.test.utils.TestUtil;
import androidx.test.ext.junit.runners.AndroidJUnit4;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.robolectric.shadows.ShadowLooper;

/** Tests for {@link DefaultDrmSessionManager} and {@link DefaultDrmSession}. */
// TODO: Test more branches:
// - Different sources for licenseServerUrl.
// - Multiple acquisitions & releases for same keys -> multiple requests.
// - Key denial.
@RunWith(AndroidJUnit4.class)
public class DefaultDrmSessionManagerTest {

  private static final UUID DRM_SCHEME_UUID =
      UUID.nameUUIDFromBytes(TestUtil.createByteArray(7, 8, 9));
  private static final ImmutableList<DrmInitData.SchemeData> DRM_SCHEME_DATAS =
      ImmutableList.of(
          new DrmInitData.SchemeData(
              DRM_SCHEME_UUID, MimeTypes.VIDEO_MP4, /* data= */ TestUtil.createByteArray(1, 2, 3)));
  private static final Format FORMAT_WITH_DRM_INIT_DATA =
      new Format.Builder().setDrmInitData(new DrmInitData(DRM_SCHEME_DATAS)).build();

  @Test(timeout = 10_000)
  public void acquireSession_triggersKeyLoadAndSessionIsOpened() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);

    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .build(/* mediaDrmCallback= */ licenseServer);
    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    assertThat(drmSession.queryKeyStatus())
        .containsExactly(FakeExoMediaDrm.KEY_STATUS_KEY, FakeExoMediaDrm.KEY_STATUS_AVAILABLE);
  }

  @Test(timeout = 10_000)
  public void keepaliveEnabled_sessionsKeptForRequestedTime() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .setSessionKeepaliveMs(10_000)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    drmSession.release(/* eventDispatcher= */ null);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    ShadowLooper.idleMainLooper(10, SECONDS);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);
  }

  @Test(timeout = 10_000)
  public void keepaliveDisabled_sessionsReleasedImmediately() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .setSessionKeepaliveMs(C.TIME_UNSET)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);
    drmSession.release(/* eventDispatcher= */ null);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);
  }

  @Test(timeout = 10_000)
  public void managerRelease_allKeepaliveSessionsImmediatelyReleased() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .setSessionKeepaliveMs(10_000)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);
    drmSession.release(/* eventDispatcher= */ null);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    drmSessionManager.release();
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);
  }

  @Test(timeout = 10_000)
  public void managerRelease_keepaliveDisabled_doesntReleaseAnySessions() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .setSessionKeepaliveMs(C.TIME_UNSET)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);

    // Release the manager, the session should still be open (though it's unusable because
    // the underlying ExoMediaDrm is released).
    drmSessionManager.release();
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
  }

  @Test(timeout = 10_000)
  public void managerRelease_mediaDrmNotReleasedUntilLastSessionReleased() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    FakeExoMediaDrm exoMediaDrm = new FakeExoMediaDrm();
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, new AppManagedProvider(exoMediaDrm))
            .setSessionKeepaliveMs(10_000)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    drmSessionManager.release();

    // The manager is now in a 'releasing' state because the session is still active - so the
    // ExoMediaDrm instance should still be active (with 1 reference held by this test, and 1 held
    // by the manager).
    assertThat(exoMediaDrm.getReferenceCount()).isEqualTo(2);

    // And re-preparing the session shouldn't acquire another reference.
    drmSessionManager.prepare();
    assertThat(exoMediaDrm.getReferenceCount()).isEqualTo(2);
    drmSessionManager.release();

    drmSession.release(/* eventDispatcher= */ null);

    // The final session has been released, so now the ExoMediaDrm should be released too.
    assertThat(exoMediaDrm.getReferenceCount()).isEqualTo(1);

    // Re-preparing the fully released manager should now acquire another ExoMediaDrm reference.
    drmSessionManager.prepare();
    assertThat(exoMediaDrm.getReferenceCount()).isEqualTo(2);
    drmSessionManager.release();

    exoMediaDrm.release();
  }

  // https://github.com/google/ExoPlayer/issues/9193
  @Test(timeout = 10_000)
  public void
      managerReleasedBeforeSession_keepaliveEnabled_managerOnlyReleasesOneKeepaliveReference()
          throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    FakeExoMediaDrm exoMediaDrm = new FakeExoMediaDrm.Builder().build();
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, new AppManagedProvider(exoMediaDrm))
            .setSessionKeepaliveMs(10_000)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);

    // Release the manager (there's still an explicit reference to the session from acquireSession).
    // This should immediately release the manager's internal keepalive session reference.
    drmSessionManager.release();
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);

    // Ensure the manager doesn't release a *second* keepalive session reference after the timer
    // expires.
    ShadowLooper.idleMainLooper(10, SECONDS);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);

    // Release the explicit session reference.
    drmSession.release(/* eventDispatcher= */ null);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);
  }

  @Test(timeout = 10_000)
  public void maxConcurrentSessionsExceeded_allKeepAliveSessionsEagerlyReleased() throws Exception {
    ImmutableList<DrmInitData.SchemeData> secondSchemeDatas =
        ImmutableList.of(DRM_SCHEME_DATAS.get(0).copyWithData(TestUtil.createByteArray(4, 5, 6)));
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS, secondSchemeDatas);
    Format secondFormatWithDrmInitData =
        new Format.Builder().setDrmInitData(new DrmInitData(secondSchemeDatas)).build();
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(
                DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm(/* maxConcurrentSessions= */ 1))
            .setSessionKeepaliveMs(10_000)
            .setMultiSession(true)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession firstDrmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(firstDrmSession);
    firstDrmSession.release(/* eventDispatcher= */ null);

    // All external references to firstDrmSession have been released, it's being kept alive by
    // drmSessionManager's internal reference.
    assertThat(firstDrmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    DrmSession secondDrmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, secondFormatWithDrmInitData));
    // The drmSessionManager had to release firstDrmSession in order to acquire secondDrmSession.
    assertThat(firstDrmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);

    waitForOpenedWithKeys(secondDrmSession);
    assertThat(secondDrmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
  }

  @Test(timeout = 10_000)
  public void maxConcurrentSessionsExceeded_allPreacquiredAndKeepaliveSessionsEagerlyReleased()
      throws Exception {
    ImmutableList<DrmInitData.SchemeData> secondSchemeDatas =
        ImmutableList.of(DRM_SCHEME_DATAS.get(0).copyWithData(TestUtil.createByteArray(4, 5, 6)));
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS, secondSchemeDatas);
    Format secondFormatWithDrmInitData =
        new Format.Builder().setDrmInitData(new DrmInitData(secondSchemeDatas)).build();
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(
                DRM_SCHEME_UUID,
                uuid -> new FakeExoMediaDrm.Builder().setMaxConcurrentSessions(1).build())
            .setSessionKeepaliveMs(10_000)
            .setMultiSession(true)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSessionReference firstDrmSessionReference =
        checkNotNull(
            drmSessionManager.preacquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    DrmSession firstDrmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(firstDrmSession);
    firstDrmSession.release(/* eventDispatcher= */ null);

    // The direct reference to firstDrmSession has been released, it's being kept alive by both
    // firstDrmSessionReference and drmSessionManager's internal reference.
    assertThat(firstDrmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    DrmSession secondDrmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, secondFormatWithDrmInitData));
    // The drmSessionManager had to release both it's internal keep-alive reference and the
    // reference represented by firstDrmSessionReference in order to acquire secondDrmSession.
    assertThat(firstDrmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);

    waitForOpenedWithKeys(secondDrmSession);
    assertThat(secondDrmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);

    // Not needed (because the manager has already released this reference) but we call it anyway
    // for completeness.
    firstDrmSessionReference.release();
    // Clean-up
    secondDrmSession.release(/* eventDispatcher= */ null);
    drmSessionManager.release();
  }

  @Test(timeout = 10_000)
  public void sessionReacquired_keepaliveTimeOutCancelled() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .setSessionKeepaliveMs(10_000)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession firstDrmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(firstDrmSession);
    firstDrmSession.release(/* eventDispatcher= */ null);

    ShadowLooper.idleMainLooper(5, SECONDS);

    // Acquire a session for the same init data 5s in to the 10s timeout (so expect the same
    // instance).
    DrmSession secondDrmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThat(secondDrmSession).isSameInstanceAs(firstDrmSession);

    // Let the timeout definitely expire, and check the session didn't get released.
    ShadowLooper.idleMainLooper(10, SECONDS);
    assertThat(secondDrmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
  }

  @Test(timeout = 10_000)
  public void preacquireSession_loadsKeysBeforeFullAcquisition() throws Exception {
    AtomicInteger keyLoadCount = new AtomicInteger(0);
    DrmSessionEventListener.EventDispatcher eventDispatcher =
        new DrmSessionEventListener.EventDispatcher();
    eventDispatcher.addEventListener(
        Util.createHandlerForCurrentLooper(),
        new DrmSessionEventListener() {
          @Override
          public void onDrmKeysLoaded(
              int windowIndex, @Nullable MediaSource.MediaPeriodId mediaPeriodId) {
            keyLoadCount.incrementAndGet();
          }
        });
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            // Disable keepalive
            .setSessionKeepaliveMs(C.TIME_UNSET)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);

    DrmSessionReference sessionReference =
        drmSessionManager.preacquireSession(eventDispatcher, FORMAT_WITH_DRM_INIT_DATA);

    // Wait for the key load event to propagate, indicating the pre-acquired session is in
    // STATE_OPENED_WITH_KEYS.
    while (keyLoadCount.get() == 0) {
      // Allow the key response to be handled.
      ShadowLooper.idleMainLooper();
    }

    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));

    // Without idling the main/playback looper, we assert the session is already in OPENED_WITH_KEYS
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    assertThat(keyLoadCount.get()).isEqualTo(1);

    // After releasing our concrete session reference, the session is held open by the pre-acquired
    // reference.
    drmSession.release(/* eventDispatcher= */ null);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);

    // Releasing the pre-acquired reference allows the session to be fully released.
    sessionReference.release();
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);
  }

  @Test(timeout = 10_000)
  public void
      preacquireSession_releaseBeforeUnderlyingAcquisitionCompletesReleasesSessionOnceAcquired()
          throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            // Disable keepalive
            .setSessionKeepaliveMs(C.TIME_UNSET)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);

    DrmSessionReference sessionReference =
        drmSessionManager.preacquireSession(/* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA);

    // Release the pre-acquired reference before the underlying session has had a chance to be
    // constructed.
    sessionReference.release();

    // Acquiring the same session triggers a second key load (because the pre-acquired session was
    // fully released).
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED);

    waitForOpenedWithKeys(drmSession);

    drmSession.release(/* eventDispatcher= */ null);
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);
  }

  @Test(timeout = 10_000)
  public void preacquireSession_releaseManagerBeforeAcquisition_acquisitionDoesntHappen()
      throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            // Disable keepalive
            .setSessionKeepaliveMs(C.TIME_UNSET)
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);

    DrmSessionReference sessionReference =
        drmSessionManager.preacquireSession(/* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA);

    // Release the manager before the underlying session has had a chance to be constructed. This
    // will release all pre-acquired sessions.
    drmSessionManager.release();

    // Allow the acquisition event to be handled on the main/playback thread.
    ShadowLooper.idleMainLooper();

    // Re-prepare the manager so we can fully acquire the same session, and check the previous
    // pre-acquisition didn't do anything.
    drmSessionManager.prepare();
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED);
    waitForOpenedWithKeys(drmSession);

    drmSession.release(/* eventDispatcher= */ null);
    // If the (still unreleased) pre-acquired session above was linked to the same underlying
    // session then the state would still be OPENED_WITH_KEYS.
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_RELEASED);

    // Release the pre-acquired session from above (this is a no-op, but we do it anyway for
    // correctness).
    sessionReference.release();
    drmSessionManager.release();
  }

  @Test(timeout = 10_000)
  public void keyRefreshEvent_triggersKeyRefresh() throws Exception {
    FakeExoMediaDrm exoMediaDrm = new FakeExoMediaDrm();
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, new AppManagedProvider(exoMediaDrm))
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);

    DefaultDrmSession drmSession =
        (DefaultDrmSession)
            checkNotNull(
                drmSessionManager.acquireSession(
                    /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);

    assertThat(licenseServer.getReceivedSchemeDatas()).hasSize(1);

    exoMediaDrm.triggerEvent(
        drmSession::hasSessionId,
        ExoMediaDrm.EVENT_KEY_REQUIRED,
        /* extra= */ 0,
        /* data= */ Util.EMPTY_BYTE_ARRAY);

    while (licenseServer.getReceivedSchemeDatas().size() == 1) {
      // Allow the key refresh event to be handled.
      ShadowLooper.idleMainLooper();
    }

    assertThat(licenseServer.getReceivedSchemeDatas()).hasSize(2);
    assertThat(ImmutableSet.copyOf(licenseServer.getReceivedSchemeDatas())).hasSize(1);

    drmSession.release(/* eventDispatcher= */ null);
    drmSessionManager.release();
    exoMediaDrm.release();
  }

  @Test(timeout = 10_000)
  public void keyRefreshEvent_whileManagerIsReleasing_triggersKeyRefresh() throws Exception {
    FakeExoMediaDrm exoMediaDrm = new FakeExoMediaDrm();
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, new AppManagedProvider(exoMediaDrm))
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);

    DefaultDrmSession drmSession =
        (DefaultDrmSession)
            checkNotNull(
                drmSessionManager.acquireSession(
                    /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    waitForOpenedWithKeys(drmSession);

    assertThat(licenseServer.getReceivedSchemeDatas()).hasSize(1);

    drmSessionManager.release();

    exoMediaDrm.triggerEvent(
        drmSession::hasSessionId,
        ExoMediaDrm.EVENT_KEY_REQUIRED,
        /* extra= */ 0,
        /* data= */ Util.EMPTY_BYTE_ARRAY);

    while (licenseServer.getReceivedSchemeDatas().size() == 1) {
      // Allow the key refresh event to be handled.
      ShadowLooper.idleMainLooper();
    }

    assertThat(licenseServer.getReceivedSchemeDatas()).hasSize(2);
    assertThat(ImmutableSet.copyOf(licenseServer.getReceivedSchemeDatas())).hasSize(1);

    drmSession.release(/* eventDispatcher= */ null);
    exoMediaDrm.release();
  }

  @Test
  public void
      deviceNotProvisioned_exceptionThrownFromOpenSession_provisioningDoneAndOpenSessionRetried() {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);

    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(
                DRM_SCHEME_UUID,
                uuid -> new FakeExoMediaDrm.Builder().setProvisionsRequired(1).build())
            .build(/* mediaDrmCallback= */ licenseServer);
    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    // Confirm that opening the session threw NotProvisionedException (otherwise state would be
    // OPENED)
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENING);
    waitForOpenedWithKeys(drmSession);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    assertThat(drmSession.queryKeyStatus())
        .containsExactly(FakeExoMediaDrm.KEY_STATUS_KEY, FakeExoMediaDrm.KEY_STATUS_AVAILABLE);
    assertThat(licenseServer.getReceivedProvisionRequests()).hasSize(1);
  }

  @Test
  public void
      deviceNotProvisioned_exceptionThrownFromGetKeyRequest_provisioningDoneAndOpenSessionRetried() {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);

    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(
                DRM_SCHEME_UUID,
                uuid ->
                    new FakeExoMediaDrm.Builder()
                        .setProvisionsRequired(1)
                        .throwNotProvisionedExceptionFromGetKeyRequest()
                        .build())
            .build(/* mediaDrmCallback= */ licenseServer);
    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED);
    waitForOpenedWithKeys(drmSession);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    assertThat(drmSession.queryKeyStatus())
        .containsExactly(FakeExoMediaDrm.KEY_STATUS_KEY, FakeExoMediaDrm.KEY_STATUS_AVAILABLE);
    assertThat(licenseServer.getReceivedProvisionRequests()).hasSize(1);
  }

  @Test
  public void deviceNotProvisioned_doubleProvisioningHandledAndOpenSessionRetried() {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);

    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(
                DRM_SCHEME_UUID,
                uuid -> new FakeExoMediaDrm.Builder().setProvisionsRequired(2).build())
            .build(/* mediaDrmCallback= */ licenseServer);
    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    // Confirm that opening the session threw NotProvisionedException (otherwise state would be
    // OPENED)
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENING);
    waitForOpenedWithKeys(drmSession);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    assertThat(drmSession.queryKeyStatus())
        .containsExactly(FakeExoMediaDrm.KEY_STATUS_KEY, FakeExoMediaDrm.KEY_STATUS_AVAILABLE);
    assertThat(licenseServer.getReceivedProvisionRequests()).hasSize(2);
  }

  @Test
  public void keyResponseIndicatesProvisioningRequired_provisioningDone() {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.requiringProvisioningThenAllowingSchemeDatas(
            DRM_SCHEME_DATAS);

    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(
                DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm.Builder().build())
            .build(/* mediaDrmCallback= */ licenseServer);
    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED);
    waitForOpenedWithKeys(drmSession);

    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
    assertThat(drmSession.queryKeyStatus())
        .containsExactly(FakeExoMediaDrm.KEY_STATUS_KEY, FakeExoMediaDrm.KEY_STATUS_AVAILABLE);
    assertThat(licenseServer.getReceivedProvisionRequests()).hasSize(1);
  }

  @Test
  public void provisioningUndoneWhileManagerIsActive_deviceReprovisioned() {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);

    FakeExoMediaDrm mediaDrm = new FakeExoMediaDrm.Builder().setProvisionsRequired(2).build();
    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, new AppManagedProvider(mediaDrm))
            .setSessionKeepaliveMs(C.TIME_UNSET)
            .build(/* mediaDrmCallback= */ licenseServer);
    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    // Confirm that opening the session threw NotProvisionedException (otherwise state would be
    // OPENED)
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENING);
    waitForOpenedWithKeys(drmSession);
    drmSession.release(/* eventDispatcher= */ null);

    mediaDrm.resetProvisioning();

    drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    // Confirm that opening the session threw NotProvisionedException (otherwise state would be
    // OPENED)
    assertThat(drmSession.getState()).isEqualTo(DrmSession.STATE_OPENING);
    waitForOpenedWithKeys(drmSession);
    assertThat(licenseServer.getReceivedProvisionRequests()).hasSize(4);
  }

  @Test
  public void managerNotPrepared_acquireSessionAndPreacquireSessionFail() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .build(/* mediaDrmCallback= */ licenseServer);

    assertThrows(
        Exception.class,
        () ->
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThrows(
        Exception.class,
        () ->
            drmSessionManager.preacquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
  }

  @Test
  public void managerReleasing_acquireSessionAndPreacquireSessionFail() throws Exception {
    FakeExoMediaDrm.LicenseServer licenseServer =
        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
    DefaultDrmSessionManager drmSessionManager =
        new DefaultDrmSessionManager.Builder()
            .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
            .build(/* mediaDrmCallback= */ licenseServer);

    drmSessionManager.prepare();
    drmSessionManager.setPlayer(/* playbackLooper= */ Looper.myLooper(), PlayerId.UNSET);
    DrmSession drmSession =
        checkNotNull(
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    drmSessionManager.release();

    // The manager's prepareCount is now zero, but the drmSession is keeping it in a 'releasing'
    // state. acquireSession and preacquireSession should still fail.
    assertThrows(
        Exception.class,
        () ->
            drmSessionManager.acquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));
    assertThrows(
        Exception.class,
        () ->
            drmSessionManager.preacquireSession(
                /* eventDispatcher= */ null, FORMAT_WITH_DRM_INIT_DATA));

    drmSession.release(/* eventDispatcher= */ null);
  }

  private static void waitForOpenedWithKeys(DrmSession drmSession) {
    while (drmSession.getState() != DrmSession.STATE_OPENED_WITH_KEYS) {
      // Check the error first, so we get a meaningful failure if there's been an error.
      assertThat(drmSession.getError()).isNull();
      assertThat(drmSession.getState()).isAnyOf(DrmSession.STATE_OPENING, DrmSession.STATE_OPENED);
      // Allow the key response to be handled.
      ShadowLooper.idleMainLooper();
    }
  }
}
