[eme] Break mojo connection during call

media/mojo/clients/mojo_cdm_unittest.cc

Index: media/mojo/clients/mojo_cdm_unittest.cc
diff --git a/media/mojo/clients/mojo_cdm_unittest.cc b/media/mojo/clients/mojo_cdm_unittest.cc
index 43c7da4e65c6546fea9a01980ef957d43f7a3bcc..4997d299fc1d4ba3e9b27ef914a77d99aa79d476 100644
--- a/media/mojo/clients/mojo_cdm_unittest.cc
+++ b/media/mojo/clients/mojo_cdm_unittest.cc
@@ -7,8 +7,10 @@
 #include "base/bind.h"
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
+#include "base/memory/ref_counted.h"
 #include "base/run_loop.h"
 #include "base/test/test_message_loop.h"
+#include "base/time/time.h"
 #include "media/base/cdm_config.h"
 #include "media/base/content_decryption_module.h"
 #include "media/base/mock_filters.h"
@@ -21,12 +23,34 @@
 #include "testing/gtest/include/gtest/gtest.h"
 
 using ::testing::_;
+using ::testing::DoAll;
+using ::testing::Invoke;
+using ::testing::NiceMock;
 using ::testing::StrictMock;
+using ::testing::WithArg;
+using ::testing::WithArgs;
+using ::testing::WithoutArgs;
 
 MATCHER(NotEmpty, "") {
   return !arg.empty();
 }
 
+ACTION_P2(CdmCreated, cdm, error_message) {
+  arg0.Run(cdm, error_message);
+}
+
+ACTION_P(ResolvePromise, cdm) {
+  cdm->ResolvePromise(arg0);
+}
+
+ACTION_P(RejectPromise, cdm) {
+  cdm->RejectPromise(arg0);
+}
+
+ACTION_P2(ResolvePromiseWithSession, cdm, session_id) {
+  cdm->ResolvePromiseWithSession(arg0, session_id);
+}
+
 namespace media {
 
 namespace {
@@ -45,7 +69,12 @@ const uint8_t kKeyId[] = {
 
 class MojoCdmTest : public ::testing::Test {
  public:
-  enum ExpectedResult { SUCCESS, CONNECTION_ERROR, FAILURE };
+  enum ExpectedResult {
+    SUCCESS,
+    CONNECTION_ERROR_BEFORE,
+    CONNECTION_ERROR_DURING,
+    FAILURE
+  };
 
   MojoCdmTest()
       : mojo_cdm_service_(base::MakeUnique<MojoCdmService>(
@@ -55,22 +84,47 @@ class MojoCdmTest : public ::testing::Test {
 
   virtual ~MojoCdmTest() {}
 
-  void Initialize(const std::string& key_system,
-                  ExpectedResult expected_result) {
+  void Initialize(ExpectedResult expected_result) {
     mojom::ContentDecryptionModulePtr remote_cdm;
     auto cdm_request = mojo::MakeRequest(&remote_cdm);
 
+    cdm_binding_.Bind(std::move(cdm_request));
+    remote_cdm_ = new NiceMock<MockCdm>();

[xhwang, 2017/01/05 06:59:53]

Can we use StrictMock?

[jrummell, 2017/01/05 23:47:22]

Yes, but then we have to ignore calls to GetCdmContext(). Done.

+
     switch (expected_result) {
       case SUCCESS:
+        // Return |remote_cdm_| when creating the CDM. SetCallbacks() is
+        // called so the MockCdm has the correct callbacks to use.
+        EXPECT_CALL(cdm_factory_, Create(_, _, _, _, _, _, _, _))
+            .WillOnce(DoAll(WithArgs<3, 4, 5, 6>(Invoke(
+                                remote_cdm_.get(), &MockCdm::SetCallbacks)),
+                            WithArg<7>(CdmCreated(remote_cdm_, ""))));

[xhwang, 2017/01/05 06:59:53]

This is pretty cool but seems a bit overkill. For the foreseeable future,
MockCdmFactory will only create MockCdm. Does it make sense to just have a
concrete MockCdmFactory implementation which will either return a MockCdm, or
nullptr, based on a flag (similar to how we construct MockCdmPromise)? Then this
code and MockCdm impl will be a bit simpler.

[jrummell, 2017/01/05 23:47:22]

Done.

+        break;
+
       case FAILURE:
-        cdm_binding_.Bind(std::move(cdm_request));
+        // Don't return a CDM.
+        EXPECT_CALL(cdm_factory_, Create(_, _, _, _, _, _, _, _))
+            .WillOnce(WithArg<7>(CdmCreated(nullptr, "CDM creation failed")));
         break;
-      case CONNECTION_ERROR:
-        cdm_request.ResetWithReason(0, "Request dropped for testing.");
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call.
+        ForceConnectionError();
+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection during the call. Also return the CDM, which
+        // should get ignored once the pipe is broken.
+        EXPECT_CALL(cdm_factory_, Create(_, _, _, _, _, _, _, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArgs<3, 4, 5, 6>(
+                    Invoke(remote_cdm_.get(), &MockCdm::SetCallbacks)),
+                WithArg<7>(CdmCreated(remote_cdm_, ""))));
         break;
     }
 
-    MojoCdm::Create(key_system, GURL(kTestSecurityOrigin), CdmConfig(),
+    MojoCdm::Create(kClearKeyKeySystem, GURL(kTestSecurityOrigin), CdmConfig(),
                     std::move(remote_cdm),
                     base::Bind(&MockCdmClient::OnSessionMessage,
                                base::Unretained(&cdm_client_)),
@@ -82,7 +136,6 @@ class MojoCdmTest : public ::testing::Test {
                                base::Unretained(&cdm_client_)),
                     base::Bind(&MojoCdmTest::OnCdmCreated,
                                base::Unretained(this), expected_result));
-
     base::RunLoop().RunUntilIdle();
   }
 
@@ -100,119 +153,505 @@ class MojoCdmTest : public ::testing::Test {
   }
 
   void ForceConnectionError() {
-    // If there is an existing session it will get closed when the connection
-    // is broken.
-    if (!session_id_.empty()) {
-      EXPECT_CALL(cdm_client_, OnSessionClosed(session_id_));
-    }
-
     cdm_binding_.CloseWithReason(2, "Test closed connection.");
-
     base::RunLoop().RunUntilIdle();
   }
 
   void SetServerCertificateAndExpect(const std::vector<uint8_t>& certificate,
                                      ExpectedResult expected_result) {
+    switch (expected_result) {
+      case SUCCESS:
+        // Resolve the promise when SetServerCertificate() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnSetServerCertificate(certificate, _))
+            .WillOnce(WithArg<1>(ResolvePromise(remote_cdm_.get())));
+        break;
+
+      case FAILURE:
+        // Reject the promise when SetServerCertificate() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnSetServerCertificate(certificate, _))
+            .WillOnce(WithArg<1>(RejectPromise(remote_cdm_.get())));
+        break;
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call, so SetServerCertificate() is
+        // never called.
+        ForceConnectionError();
+        EXPECT_CALL(*remote_cdm_.get(), OnSetServerCertificate(_, _)).Times(0);

[xhwang, 2017/01/05 06:59:53]

If you use a StrongMock, you don't need Times(0) because uninteresting call
would also fail the test.

[jrummell, 2017/01/05 23:47:22]

Done.

+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection before resolving the promise. Note that the
+        // resolved promise never gets sent across the broken pipe.
+        EXPECT_CALL(*remote_cdm_.get(), OnSetServerCertificate(certificate, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArg<1>(ResolvePromise(remote_cdm_.get()))));

[xhwang, 2017/01/05 06:59:53]

nit: Can we try not resolving the promise here? I guess it's important that even
if the remote CDM resolves the promise, MojoCdm will still get a rejected
promise. But it's equally or more important that even if the remote CDM does
nothing, e.g. the remote CDM process crashed, MojoCdm will also get a rejected
promise.

[jrummell, 2017/01/05 23:47:22]

Rejecting or resolving the promise doesn't get passed back to MojoCdm. This is
simply done so that the promise doesn't complain that it's never fulfilled.

+        break;
+    }
+
     mojo_cdm_->SetServerCertificate(
         certificate,
         base::MakeUnique<MockCdmPromise>(expected_result == SUCCESS));
-
     base::RunLoop().RunUntilIdle();
   }
 
-  void CreateSessionAndExpect(EmeInitDataType data_type,
-                              const std::vector<uint8_t>& key_id,
+  void CreateSessionAndExpect(const std::string& session_id,
                               ExpectedResult expected_result) {
-    if (expected_result == SUCCESS) {
-      EXPECT_CALL(cdm_client_, OnSessionMessage(NotEmpty(), _, _));
+    // Specify parameters needed to call CreateSessionAndGenerateRequest() in
+    // order to verify that the data is passed properly.
+    const CdmSessionType session_type = CdmSessionType::TEMPORARY_SESSION;
+    const EmeInitDataType data_type = EmeInitDataType::WEBM;
+    const std::vector<uint8_t> key_id(kKeyId, kKeyId + arraysize(kKeyId));
+    std::string created_session_id;
+
+    switch (expected_result) {
+      case SUCCESS:
+        // Resolve the promise when CreateSessionAndGenerateRequest() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnCreateSessionAndGenerateRequest(
+                                            session_type, data_type, key_id, _))
+            .WillOnce(WithArg<3>(
+                ResolvePromiseWithSession(remote_cdm_.get(), session_id)));
+        break;
+
+      case FAILURE:
+        // Reject the promise when CreateSessionAndGenerateRequest() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnCreateSessionAndGenerateRequest(
+                                            session_type, data_type, key_id, _))
+            .WillOnce(WithArg<3>(RejectPromise(remote_cdm_.get())));
+        break;
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call, so
+        // CreateSessionAndGenerateRequest() is never called.
+        ForceConnectionError();
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnCreateSessionAndGenerateRequest(_, _, _, _))
+            .Times(0);
+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection before resolving the promise. Note that the
+        // resolved promise never gets sent across the broken pipe.
+        EXPECT_CALL(*remote_cdm_.get(), OnCreateSessionAndGenerateRequest(
+                                            session_type, data_type, key_id, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArg<3>(
+                    ResolvePromiseWithSession(remote_cdm_.get(), session_id))));
+        break;
     }

[xhwang, 2017/01/05 06:59:53]

There are a lot of duplication on this switch-case block between this and other
*AndExpect() helper functions. Is it possible to do something like:

if (expected_result == CONNECTION_ERROR_BEFORE) {
  ForceConnectionError();
} else {
  EXPECT_CALL(*remote_cdm_.get(), OnCreateSessionAndGenerateRequest(
                                             session_type, data_type, key_id,
_))
      .WillOnce(WithArg<3>(ProcessPromise(..., expected_result)));
}

Then in ProcessPromise, we can check |expected_result| and decide whether to
resolve or reject the promise, or trigger a connection error.

[jrummell, 2017/01/05 23:47:22]

ProcessPromise is a bit of work, but done.

 
+    // Note that although it's called CreateSessionAndGenerateRequest, no
+    // request is generated.
     mojo_cdm_->CreateSessionAndGenerateRequest(
-        CdmSessionType::TEMPORARY_SESSION, data_type, key_id,
+        session_type, data_type, key_id,
         base::MakeUnique<MockCdmSessionPromise>(expected_result == SUCCESS,
-                                                &session_id_));
-
+                                                &created_session_id));
     base::RunLoop().RunUntilIdle();
+
+    // If the session was "created" ...
+    if (expected_result == SUCCESS) {
+      // Returned session ID must match the session ID provided.
+      EXPECT_EQ(session_id, created_session_id);
+
+      // MojoCdm expects the session to be closed, so invoke SessionClosedCB
+      // to "close" it.
+      EXPECT_CALL(cdm_client_, OnSessionClosed(session_id));
+      remote_cdm_->CallSessionClosedCB(session_id);
+      base::RunLoop().RunUntilIdle();
+    }
   }
 
-  void CloseSessionAndExpect(ExpectedResult expected_result) {
-    DCHECK(!session_id_.empty()) << "CloseSessionAndExpect() must be called "
-                                    "after a successful "
-                                    "CreateSessionAndExpect()";
+  void LoadSessionAndExpect(const std::string& session_id,
+                            ExpectedResult expected_result) {
+    const CdmSessionType session_type =
+        CdmSessionType::PERSISTENT_LICENSE_SESSION;
+    std::string loaded_session_id;
+
+    switch (expected_result) {
+      case SUCCESS:
+        // Resolve the promise when LoadSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnLoadSession(session_type, session_id, _))
+            .WillOnce(WithArgs<2, 1>(Invoke(
+                remote_cdm_.get(), &MockCdm::ResolvePromiseWithSession)));
+        break;
+
+      case FAILURE:
+        // Reject the promise when LoadSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnLoadSession(session_type, session_id, _))
+            .WillOnce(WithArg<2>(RejectPromise(remote_cdm_.get())));
+        break;
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call, so LoadSession() is never
+        // called.
+        ForceConnectionError();
+        EXPECT_CALL(*remote_cdm_.get(), OnLoadSession(_, _, _)).Times(0);
+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection before resolving the promise. Note that the
+        // resolved promise never gets sent across the broken pipe.
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnLoadSession(session_type, session_id, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArgs<2, 1>(Invoke(remote_cdm_.get(),
+                                      &MockCdm::ResolvePromiseWithSession))));
+        break;
+    }
+
+    mojo_cdm_->LoadSession(session_type, session_id,
+                           base::MakeUnique<MockCdmSessionPromise>(
+                               expected_result == SUCCESS, &loaded_session_id));
+    base::RunLoop().RunUntilIdle();
 
+    // If the session was "loaded" ...
     if (expected_result == SUCCESS) {
-      EXPECT_CALL(cdm_client_, OnSessionClosed(session_id_));
+      // Returned session ID must match the session ID provided.
+      EXPECT_EQ(session_id, loaded_session_id);
+
+      // MojoCdm expects the session to be closed, so invoke SessionClosedCB
+      // to "close" it.
+      EXPECT_CALL(cdm_client_, OnSessionClosed(session_id));
+      remote_cdm_->CallSessionClosedCB(session_id);
+      base::RunLoop().RunUntilIdle();
     }
+  }
 
-    mojo_cdm_->CloseSession(session_id_, base::MakeUnique<MockCdmPromise>(
-                                             expected_result == SUCCESS));
+  void UpdateSessionAndExpect(const std::string& session_id,
+                              ExpectedResult expected_result) {
+    // Specify parameters to UpdateSession() in order to verify that
+    // the data is passed properly.
+    const std::vector<uint8_t> response = {1, 2, 3, 4, 5, 6};
+
+    switch (expected_result) {
+      case SUCCESS:
+        // Resolve the promise when UpdateSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnUpdateSession(session_id, response, _))
+            .WillOnce(WithArg<2>(ResolvePromise(remote_cdm_.get())));
+        break;
+
+      case FAILURE:
+        // Reject the promise when UpdateSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnUpdateSession(session_id, response, _))
+            .WillOnce(WithArg<2>(RejectPromise(remote_cdm_.get())));
+        break;
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call, so UpdateSession() is never
+        // called.
+        ForceConnectionError();
+        EXPECT_CALL(*remote_cdm_.get(), OnUpdateSession(_, _, _)).Times(0);
+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection before resolving the promise. Note that the
+        // resolved promise never gets sent across the broken pipe.
+        EXPECT_CALL(*remote_cdm_.get(),
+                    OnUpdateSession(session_id, response, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArg<2>(ResolvePromise(remote_cdm_.get()))));
+        break;
+    }
 
+    mojo_cdm_->UpdateSession(
+        session_id, response,
+        base::MakeUnique<MockCdmPromise>(expected_result == SUCCESS));
+    base::RunLoop().RunUntilIdle();
+  }
+
+  void CloseSessionAndExpect(const std::string& session_id,
+                             ExpectedResult expected_result) {
+    switch (expected_result) {
+      case SUCCESS:
+        // Resolve the promise when CloseSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnCloseSession(session_id, _))
+            .WillOnce(WithArg<1>(ResolvePromise(remote_cdm_.get())));
+        break;
+
+      case FAILURE:
+        // Reject the promise when CloseSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnCloseSession(session_id, _))
+            .WillOnce(WithArg<1>(RejectPromise(remote_cdm_.get())));
+        break;
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call, so CloseSession() is never
+        // called.
+        ForceConnectionError();
+        EXPECT_CALL(*remote_cdm_.get(), OnCloseSession(_, _)).Times(0);
+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection before resolving the promise. Note that the
+        // resolved promise never gets sent across the broken pipe.
+        EXPECT_CALL(*remote_cdm_.get(), OnCloseSession(session_id, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArg<1>(ResolvePromise(remote_cdm_.get()))));
+        break;
+    }
+
+    mojo_cdm_->CloseSession(session_id, base::MakeUnique<MockCdmPromise>(
+                                            expected_result == SUCCESS));
+    base::RunLoop().RunUntilIdle();
+  }
+
+  void RemoveSessionAndExpect(const std::string& session_id,
+                              ExpectedResult expected_result) {
+    switch (expected_result) {
+      case SUCCESS:
+        // Resolve the promise when RemoveSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnRemoveSession(session_id, _))
+            .WillOnce(WithArg<1>(ResolvePromise(remote_cdm_.get())));
+        break;
+
+      case FAILURE:
+        // Reject the promise when RemoveSession() is called.
+        EXPECT_CALL(*remote_cdm_.get(), OnRemoveSession(session_id, _))
+            .WillOnce(WithArg<1>(RejectPromise(remote_cdm_.get())));
+        break;
+
+      case CONNECTION_ERROR_BEFORE:
+        // Break the connection before the call, so RemoveSession() never
+        // called.
+        ForceConnectionError();
+        EXPECT_CALL(*remote_cdm_.get(), OnRemoveSession(_, _)).Times(0);
+        break;
+
+      case CONNECTION_ERROR_DURING:
+        // Break the connection before resolving the promise. Note that the
+        // resolved promise never gets sent across the broken pipe.
+        EXPECT_CALL(*remote_cdm_.get(), OnRemoveSession(session_id, _))
+            .WillOnce(DoAll(
+                WithoutArgs(Invoke(this, &MojoCdmTest::ForceConnectionError)),
+                WithArg<1>(ResolvePromise(remote_cdm_.get()))));
+        break;
+    }
+
+    mojo_cdm_->RemoveSession(session_id, base::MakeUnique<MockCdmPromise>(
+                                             expected_result == SUCCESS));
     base::RunLoop().RunUntilIdle();
   }
 
   // Fixture members.
   base::TestMessageLoop message_loop_;
 
+  // |remote_cdm_| represents the CDM at the end of the mojo message pipe.
+  // |cdm_factory_| will return |remote_cdm_| when the CDM is needed.
+  scoped_refptr<NiceMock<MockCdm>> remote_cdm_;
+  MockCdmFactory cdm_factory_;
+
   MojoCdmServiceContext mojo_cdm_service_context_;
   StrictMock<MockCdmClient> cdm_client_;
 
-  // TODO(jrummell): Use a MockCdmFactory to create a MockCdm here for more test
-  // coverage.
-  DefaultCdmFactory cdm_factory_;
-
   std::unique_ptr<MojoCdmService> mojo_cdm_service_;
   mojo::Binding<mojom::ContentDecryptionModule> cdm_binding_;
   scoped_refptr<ContentDecryptionModule> mojo_cdm_;
 
-  // |session_id_| is the latest successful result of calling CreateSession().
-  std::string session_id_;
-
  private:
   DISALLOW_COPY_AND_ASSIGN(MojoCdmTest);
 };
 
 TEST_F(MojoCdmTest, Create_Success) {
-  Initialize(kClearKeyKeySystem, SUCCESS);
+  Initialize(SUCCESS);
 }
 
-TEST_F(MojoCdmTest, Create_ConnectionError) {
-  Initialize(kClearKeyKeySystem, CONNECTION_ERROR);
+TEST_F(MojoCdmTest, Create_Failure) {
+  Initialize(FAILURE);
 }
 
-TEST_F(MojoCdmTest, Create_Failure) {
-  // This fails as DefaultCdmFactory only supports Clear Key.
-  Initialize("org.random.cdm", FAILURE);
+TEST_F(MojoCdmTest, Create_ConnectionErrorBefore) {
+  Initialize(CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, Create_ConnectionErrorDuring) {
+  Initialize(CONNECTION_ERROR_DURING);
 }
 
-TEST_F(MojoCdmTest, SetServerCertificate_AfterConnectionError) {
-  Initialize(kClearKeyKeySystem, SUCCESS);
-  ForceConnectionError();
-  SetServerCertificateAndExpect({0, 1, 2}, FAILURE);
+TEST_F(MojoCdmTest, SetServerCertificate_Success) {
+  const std::vector<uint8_t> certificate = {0, 1, 2};
+  Initialize(SUCCESS);
+  SetServerCertificateAndExpect(certificate, SUCCESS);
 }
 
-TEST_F(MojoCdmTest, CreateSessionAndGenerateRequest_AfterConnectionError) {
-  std::vector<uint8_t> key_id(kKeyId, kKeyId + arraysize(kKeyId));
+TEST_F(MojoCdmTest, SetServerCertificate_Failure) {
+  const std::vector<uint8_t> certificate = {1, 2, 3, 4, 5};
+  Initialize(SUCCESS);
+  SetServerCertificateAndExpect(certificate, FAILURE);
+}
 
-  Initialize(kClearKeyKeySystem, SUCCESS);
-  ForceConnectionError();
-  CreateSessionAndExpect(EmeInitDataType::WEBM, key_id, FAILURE);
+TEST_F(MojoCdmTest, SetServerCertificate_ConnectionErrorBefore) {
+  const std::vector<uint8_t> certificate = {3, 4};
+  Initialize(SUCCESS);
+  SetServerCertificateAndExpect(certificate, CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, SetServerCertificate_ConnectionErrorDuring) {
+  const std::vector<uint8_t> certificate = {10, 11, 12};
+  Initialize(SUCCESS);
+  SetServerCertificateAndExpect(certificate, CONNECTION_ERROR_DURING);
+}
+
+TEST_F(MojoCdmTest, CreateSession_Success) {
+  const std::string session_id = "create1";
+  Initialize(SUCCESS);
+  CreateSessionAndExpect("create1", SUCCESS);
+}
+
+TEST_F(MojoCdmTest, CreateSession_Failure) {
+  const std::string session_id = "create2";
+  Initialize(SUCCESS);
+  CreateSessionAndExpect(session_id, FAILURE);
+}
+
+TEST_F(MojoCdmTest, CreateSession_ConnectionErrorBefore) {
+  const std::string session_id = "create3";
+  Initialize(SUCCESS);
+  CreateSessionAndExpect(session_id, CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, CreateSession_ConnectionErrorDuring) {
+  const std::string session_id = "create4";
+  Initialize(SUCCESS);
+  CreateSessionAndExpect(session_id, CONNECTION_ERROR_DURING);
+}
+
+TEST_F(MojoCdmTest, LoadSession_Success) {
+  const std::string session_id = "load1";
+  Initialize(SUCCESS);
+  LoadSessionAndExpect(session_id, SUCCESS);
+}
+
+TEST_F(MojoCdmTest, LoadSession_Failure) {
+  const std::string session_id = "load2";
+  Initialize(SUCCESS);
+  LoadSessionAndExpect(session_id, FAILURE);
+}
+
+TEST_F(MojoCdmTest, LoadSession_ConnectionErrorBefore) {
+  const std::string session_id = "load3";
+  Initialize(SUCCESS);
+  LoadSessionAndExpect(session_id, CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, LoadSession_ConnectionErrorDuring) {
+  const std::string session_id = "load4";
+  Initialize(SUCCESS);
+  LoadSessionAndExpect(session_id, CONNECTION_ERROR_DURING);
+}
+
+TEST_F(MojoCdmTest, UpdateSession_Success) {
+  const std::string session_id = "update1";
+  Initialize(SUCCESS);
+  UpdateSessionAndExpect(session_id, SUCCESS);
+}
+
+TEST_F(MojoCdmTest, UpdateSession_Failure) {
+  const std::string session_id = "update2";
+  Initialize(SUCCESS);
+  UpdateSessionAndExpect(session_id, FAILURE);
+}
+
+TEST_F(MojoCdmTest, UpdateSession_ConnectionErrorBefore) {
+  const std::string session_id = "update3";
+  Initialize(SUCCESS);
+  UpdateSessionAndExpect(session_id, CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, UpdateSession_ConnectionErrorDuring) {
+  const std::string session_id = "update4";
+  Initialize(SUCCESS);
+  UpdateSessionAndExpect(session_id, CONNECTION_ERROR_DURING);
 }
 
 TEST_F(MojoCdmTest, CloseSession_Success) {
-  std::vector<uint8_t> key_id(kKeyId, kKeyId + arraysize(kKeyId));
+  const std::string session_id = "close1";
+  Initialize(SUCCESS);
+  CloseSessionAndExpect(session_id, SUCCESS);
+}
+
+TEST_F(MojoCdmTest, CloseSession_Failure) {
+  const std::string session_id = "close2";
+  Initialize(SUCCESS);
+  CloseSessionAndExpect(session_id, FAILURE);
+}
+
+TEST_F(MojoCdmTest, CloseSession_ConnectionErrorBefore) {
+  const std::string session_id = "close3";
+  Initialize(SUCCESS);
+  CloseSessionAndExpect(session_id, CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, CloseSession_ConnectionErrorDuring) {
+  const std::string session_id = "close4";
+  Initialize(SUCCESS);
+  CloseSessionAndExpect(session_id, CONNECTION_ERROR_DURING);
+}
 
-  Initialize(kClearKeyKeySystem, SUCCESS);
-  CreateSessionAndExpect(EmeInitDataType::WEBM, key_id, SUCCESS);
-  CloseSessionAndExpect(SUCCESS);
+TEST_F(MojoCdmTest, RemoveSession_Success) {
+  const std::string session_id = "remove1";
+  Initialize(SUCCESS);
+  RemoveSessionAndExpect(session_id, SUCCESS);
 }
 
-TEST_F(MojoCdmTest, CloseSession_AfterConnectionError) {
-  std::vector<uint8_t> key_id(kKeyId, kKeyId + arraysize(kKeyId));
+TEST_F(MojoCdmTest, RemoveSession_Failure) {
+  const std::string session_id = "remove2";
+  Initialize(SUCCESS);
+  RemoveSessionAndExpect(session_id, FAILURE);
+}
+
+TEST_F(MojoCdmTest, RemoveSession_ConnectionErrorBefore) {
+  const std::string session_id = "remove3";
+  Initialize(SUCCESS);
+  RemoveSessionAndExpect(session_id, CONNECTION_ERROR_BEFORE);
+}
+
+TEST_F(MojoCdmTest, RemoveSession_ConnectionErrorDuring) {
+  const std::string session_id = "remove4";
+  Initialize(SUCCESS);
+  RemoveSessionAndExpect(session_id, CONNECTION_ERROR_DURING);
+}
+
+// Note that MojoCdm requires a session to exist when SessionClosedCB is called,
+// so it is currently tested in the success cases for CreateSession/LoadSession.
+
+TEST_F(MojoCdmTest, SessionMessageCB_Success) {
+  const std::string session_id = "message";
+  const ContentDecryptionModule::MessageType message_type =
+      ContentDecryptionModule::LICENSE_REQUEST;
+  const std::vector<uint8_t> message = {0, 1, 2};
+  Initialize(SUCCESS);
+  EXPECT_CALL(cdm_client_, OnSessionMessage(session_id, message_type, message));
+  remote_cdm_->CallSessionMessageCB(session_id, message_type, message);
+  base::RunLoop().RunUntilIdle();
+}
+
+TEST_F(MojoCdmTest, SessionExpirationChangeCB_Success) {
+  const std::string session_id = "expiration";
+  const base::Time time = base::Time::Now();
+  Initialize(SUCCESS);
+  EXPECT_CALL(cdm_client_, OnSessionExpirationUpdate(session_id, time));
+  remote_cdm_->CallSessionExpirationUpdateCB(session_id, time);
+  base::RunLoop().RunUntilIdle();
+}
 
-  Initialize(kClearKeyKeySystem, SUCCESS);
-  CreateSessionAndExpect(EmeInitDataType::WEBM, key_id, SUCCESS);
-  ForceConnectionError();
-  CloseSessionAndExpect(FAILURE);
+TEST_F(MojoCdmTest, SessionKeysChangeCB_Success) {
+  const std::string session_id = "change";
+  bool has_additional_usable_key = true;
+  CdmKeysInfo keys_info;
+  Initialize(SUCCESS);
+  EXPECT_CALL(cdm_client_,
+              OnSessionKeysChangeCalled(session_id, has_additional_usable_key));
+  remote_cdm_->CallSessionKeysChangeCB(session_id, has_additional_usable_key,
+                                       std::move(keys_info));
+  base::RunLoop().RunUntilIdle();
 }
 
 }  // namespace media