Clean up Assignment Create in BlimpClientContextImpl.

blimp/client/core/context/blimp_client_context_impl.cc

Index: blimp/client/core/context/blimp_client_context_impl.cc
diff --git a/blimp/client/core/context/blimp_client_context_impl.cc b/blimp/client/core/context/blimp_client_context_impl.cc
index b8278ecfe3b0628a4417df974ec42df1da24f9e2..f1a8eeb1471f2e861d440d0393b808d3e6b7d49e 100644
--- a/blimp/client/core/context/blimp_client_context_impl.cc
+++ b/blimp/client/core/context/blimp_client_context_impl.cc
@@ -138,9 +138,16 @@ std::unique_ptr<BlimpContents> BlimpClientContextImpl::CreateBlimpContents(
 }
 
 void BlimpClientContextImpl::Connect() {
-  // Start Blimp authentication flow. The OAuth2 token will be used in
-  // assignment source.
-  GetIdentitySource()->Connect();
+  if (!assignment_fetcher_) {
+    assignment_fetcher_ = base::MakeUnique<AssignmentFetcher>(
+        delegate_->CreateIdentityProvider(),
+        base::Bind(&BlimpClientContextDelegate::OnAuthenticationError,
+                   base::Unretained(delegate_)),
+        GetAssignerURL(), io_thread_task_runner_, file_thread_task_runner_);
+  }
+  assignment_fetcher_->Fetch(
+      base::Bind(&BlimpClientContextImpl::OnAssignmentReceived,
+                 weak_factory_.GetWeakPtr()));
 }
 
 void BlimpClientContextImpl::ConnectWithAssignment(
@@ -156,35 +163,18 @@ BlimpClientContextImpl::CreateFeedbackData() {
   return CreateBlimpFeedbackData(blimp_contents_manager_.get());
 }
 
-void BlimpClientContextImpl::OnAuthTokenReceived(
-    const std::string& client_auth_token) {
-  if (!assignment_source_) {
-    assignment_source_.reset(new AssignmentSource(
-        GetAssignerURL(), io_thread_task_runner_, file_thread_task_runner_));
-  }
-
-  VLOG(1) << "Trying to get assignment.";
-  assignment_source_->GetAssignment(
-      client_auth_token,
-      base::Bind(&BlimpClientContextImpl::OnAssignmentReceived,
-                 weak_factory_.GetWeakPtr()));
-}
-
-GURL BlimpClientContextImpl::GetAssignerURL() {
-  return GURL(kDefaultAssignerUrl);
-}
-
 IdentitySource* BlimpClientContextImpl::GetIdentitySource() {
-  if (!identity_source_) {
-    CreateIdentitySource();
-  }
-  return identity_source_.get();
+  return assignment_fetcher_->GetIdentitySource();
 }
 
 ConnectionStatus* BlimpClientContextImpl::GetConnectionStatus() {
   return &connection_status_;
 }
 
+GURL BlimpClientContextImpl::GetAssignerURL() {
+  return GURL(kDefaultAssignerUrl);
+}
+
 void BlimpClientContextImpl::OnAssignmentReceived(
     AssignmentRequestResult result,
     const Assignment& assignment) {
@@ -246,12 +236,6 @@ void BlimpClientContextImpl::DropConnection() {
       FROM_HERE, base::Bind(&DropConnectionOnIOThread, net_components_.get()));
 }
 
-void BlimpClientContextImpl::CreateIdentitySource() {
-  identity_source_ = base::MakeUnique<IdentitySource>(
-      delegate_, base::Bind(&BlimpClientContextImpl::OnAuthTokenReceived,
-                            base::Unretained(this)));
-}
-
 void BlimpClientContextImpl::OnImageDecodeError() {
   // Currently we just drop the connection on image decoding error.
   DropConnection();