From 84cf63a72f90a2a841fa5e233faf2e6935a77232 Mon Sep 17 00:00:00 2001
From: ibaker <ibaker@google.com>
Date: Thu, 26 Aug 2021 16:45:03 +0100
Subject: [PATCH] Test DefaultDrmSession provisioning is requested by
 provideKeyResponse

Follow-up to
<unknown commit>

PiperOrigin-RevId: 393132950
---
 .../drm/DefaultDrmSessionManagerTest.java     | 27 +++++++++++++++
 .../exoplayer2/testutil/FakeExoMediaDrm.java  | 33 +++++++++++++++++--
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/library/core/src/test/java/com/google/android/exoplayer2/drm/DefaultDrmSessionManagerTest.java b/library/core/src/test/java/com/google/android/exoplayer2/drm/DefaultDrmSessionManagerTest.java
index d5bfc5f663..e5375012f7 100644
--- a/library/core/src/test/java/com/google/android/exoplayer2/drm/DefaultDrmSessionManagerTest.java
+++ b/library/core/src/test/java/com/google/android/exoplayer2/drm/DefaultDrmSessionManagerTest.java
@@ -665,6 +665,33 @@ public class DefaultDrmSessionManagerTest {
     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();
+    DrmSession drmSession =
+        checkNotNull(
+            drmSessionManager.acquireSession(
+                /* playbackLooper= */ checkNotNull(Looper.myLooper()),
+                /* 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 =
diff --git a/testutils/src/main/java/com/google/android/exoplayer2/testutil/FakeExoMediaDrm.java b/testutils/src/main/java/com/google/android/exoplayer2/testutil/FakeExoMediaDrm.java
index 9e5d1498b6..a369984242 100644
--- a/testutils/src/main/java/com/google/android/exoplayer2/testutil/FakeExoMediaDrm.java
+++ b/testutils/src/main/java/com/google/android/exoplayer2/testutil/FakeExoMediaDrm.java
@@ -138,6 +138,8 @@ public final class FakeExoMediaDrm implements ExoMediaDrm {
 
   private static final ImmutableList<Byte> VALID_KEY_RESPONSE = TestUtil.createByteList(1, 2, 3);
   private static final ImmutableList<Byte> KEY_DENIED_RESPONSE = TestUtil.createByteList(9, 8, 7);
+  private static final ImmutableList<Byte> PROVISIONING_REQUIRED_RESPONSE =
+      TestUtil.createByteList(4, 5, 6);
 
   private final int provisionsRequired;
   private final int maxConcurrentSessions;
@@ -258,7 +260,6 @@ public final class FakeExoMediaDrm implements ExoMediaDrm {
     return new KeyRequest(requestData.toByteArray(), /* licenseServerUrl= */ "", requestType);
   }
 
-  @Nullable
   @Override
   public byte[] provideKeyResponse(byte[] scope, byte[] response)
       throws NotProvisionedException, DeniedByServerException {
@@ -268,6 +269,8 @@ public final class FakeExoMediaDrm implements ExoMediaDrm {
       sessionIdsWithValidKeys.add(Bytes.asList(scope));
     } else if (responseAsList.equals(KEY_DENIED_RESPONSE)) {
       throw new DeniedByServerException("Key request denied");
+    } else if (responseAsList.equals(PROVISIONING_REQUIRED_RESPONSE)) {
+      throw new NotProvisionedException("Provisioning required");
     }
     return Util.EMPTY_BYTE_ARRAY;
   }
@@ -417,6 +420,8 @@ public final class FakeExoMediaDrm implements ExoMediaDrm {
     private final List<ImmutableList<Byte>> receivedProvisionRequests;
     private final List<ImmutableList<DrmInitData.SchemeData>> receivedSchemeDatas;
 
+    private boolean nextResponseIndicatesProvisioningRequired;
+
     @SafeVarargs
     public static LicenseServer allowingSchemeDatas(List<DrmInitData.SchemeData>... schemeDatas) {
       ImmutableSet.Builder<ImmutableList<DrmInitData.SchemeData>> schemeDatasBuilder =
@@ -427,6 +432,19 @@ public final class FakeExoMediaDrm implements ExoMediaDrm {
       return new LicenseServer(schemeDatasBuilder.build());
     }
 
+    @SafeVarargs
+    public static LicenseServer requiringProvisioningThenAllowingSchemeDatas(
+        List<DrmInitData.SchemeData>... schemeDatas) {
+      ImmutableSet.Builder<ImmutableList<DrmInitData.SchemeData>> schemeDatasBuilder =
+          ImmutableSet.builder();
+      for (List<DrmInitData.SchemeData> schemeData : schemeDatas) {
+        schemeDatasBuilder.add(ImmutableList.copyOf(schemeData));
+      }
+      LicenseServer licenseServer = new LicenseServer(schemeDatasBuilder.build());
+      licenseServer.nextResponseIndicatesProvisioningRequired = true;
+      return licenseServer;
+    }
+
     private LicenseServer(ImmutableSet<ImmutableList<DrmInitData.SchemeData>> allowedSchemeDatas) {
       this.allowedSchemeDatas = allowedSchemeDatas;
 
@@ -459,8 +477,17 @@ public final class FakeExoMediaDrm implements ExoMediaDrm {
       ImmutableList<DrmInitData.SchemeData> schemeDatas =
           KeyRequestData.fromByteArray(request.getData()).schemeDatas;
       receivedSchemeDatas.add(schemeDatas);
-      return Bytes.toArray(
-          allowedSchemeDatas.contains(schemeDatas) ? VALID_KEY_RESPONSE : KEY_DENIED_RESPONSE);
+
+      ImmutableList<Byte> response;
+      if (nextResponseIndicatesProvisioningRequired) {
+        nextResponseIndicatesProvisioningRequired = false;
+        response = PROVISIONING_REQUIRED_RESPONSE;
+      } else if (allowedSchemeDatas.contains(schemeDatas)) {
+        response = VALID_KEY_RESPONSE;
+      } else {
+        response = KEY_DENIED_RESPONSE;
+      }
+      return Bytes.toArray(response);
     }
   }