1 files changed, 7 insertions, 0 deletions
diff --git a/mojo/shell/public/cpp/lib/shell_connection.cc b/mojo/shell/public/cpp/lib/shell_connection.cc
index 7c8b032..d1cba7e 100644
--- a/mojo/shell/public/cpp/lib/shell_connection.cc
+++ b/mojo/shell/public/cpp/lib/shell_connection.cc
@@ -32,6 +32,10 @@ ShellConnection::ShellConnection(mojo::ShellClient* client,
 
 ShellConnection::~ShellConnection() {}
 
+void ShellConnection::set_initialize_handler(const base::Closure& callback) {
+  initialize_handler_ = callback;
+}
+
 void ShellConnection::SetAppTestConnectorForTesting(
     shell::mojom::ConnectorPtr connector) {
   pending_connector_request_ = nullptr;
@@ -44,6 +48,9 @@ void ShellConnection::SetAppTestConnectorForTesting(
 void ShellConnection::Initialize(shell::mojom::IdentityPtr identity,
                                  uint32_t id,
                                  const InitializeCallback& callback) {
+  if (!initialize_handler_.is_null())
+    initialize_handler_.Run();
+
   callback.Run(std::move(pending_connector_request_));
 
   DCHECK(binding_.is_bound());