Add token serialization to token_service.

The token service now supports serialization of tokens into the web database. The token service doesn't actually know the email of the user though, and it's assumed that whoever is using the token service does. Not sure whether that's a bad choice or not, open to suggestions.

This CL adds a new DB table.

Design comments / corrections are welcome. Testing in Chromium OS is in progress and not yet done. This CL will not be submitted prior to more Chromium OS testing.

BUG=47092,47093
TEST=Unit tests included

Review URL: http://codereview.chromium.org/3061025

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@55835 0039d316-1c4b-4281-b951-d872f2087c98

3 files changed, 467 insertions, 58 deletions

diff --git a/chrome/browser/net/gaia/token_service.cc b/chrome/browser/net/gaia/token_service.cc
index 0a87dfe..eaa7a1ae 100644
--- a/chrome/browser/net/gaia/token_service.cc
+++ b/chrome/browser/net/gaia/token_service.cc
@@ -5,26 +5,86 @@
 #include "chrome/browser/net/gaia/token_service.h"
 
 #include "base/string_util.h"
+#include "chrome/browser/profile.h"
 #include "chrome/common/net/gaia/gaia_authenticator2.h"
 #include "chrome/common/net/gaia/gaia_constants.h"
 #include "chrome/common/notification_service.h"
 
-void TokenService::Initialize(
-    const char* const source,
-    URLRequestContextGetter* getter,
-    const GaiaAuthConsumer::ClientLoginResult& credentials) {
+// Unfortunately kNumServices must be defined in the .h.
+const char* TokenService::kServices[] = {GaiaConstants::kSyncService,
+                                         GaiaConstants::kTalkService};
+TokenService::TokenService()
+    : token_loading_query_(NULL) {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+}
 
-  credentials_ = credentials;
+TokenService::~TokenService() {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  ResetCredentialsInMemory();
+}
+
+void TokenService::Initialize(const char* const source,
+                              Profile* profile) {
+
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+
+  getter_ = profile->GetRequestContext();
+  // Since the user can create a bookmark in incognito, sync may be running.
+  // Thus we have to go for explicit access.
+  web_data_service_ = profile->GetWebDataService(Profile::EXPLICIT_ACCESS);
   source_ = std::string(source);
-  sync_token_fetcher_.reset(new GaiaAuthenticator2(this, source_, getter));
-  talk_token_fetcher_.reset(new GaiaAuthenticator2(this, source_, getter));
 }
 
-bool TokenService::AreCredentialsValid() const {
+void TokenService::ResetCredentialsInMemory() {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+
+  // Terminate any running fetchers. Callbacks will not return.
+  for (int i = 0; i < kNumServices; i++) {
+    fetchers_[i].reset();
+  }
+
+  // Cancel pending loads. Callbacks will not return.
+  if (token_loading_query_) {
+    web_data_service_->CancelRequest(token_loading_query_);
+    token_loading_query_ = NULL;
+  }
+
+  token_map_.clear();
+  credentials_ = GaiaAuthConsumer::ClientLoginResult();
+}
+
+void TokenService::UpdateCredentials(
+    const GaiaAuthConsumer::ClientLoginResult& credentials) {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  credentials_ = credentials;
+
+  // Cancels any currently running requests.
+  for (int i = 0; i < kNumServices; i++) {
+    fetchers_[i].reset(new GaiaAuthenticator2(this, source_, getter_));
+  }
+}
+
+void TokenService::LoadTokensFromDB() {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  token_loading_query_ = web_data_service_->GetAllTokens(this);
+}
+
+void TokenService::SaveAuthTokenToDB(const std::string& service,
+                                     const std::string& auth_token) {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  web_data_service_->SetTokenForService(service, auth_token);
+}
+
+void TokenService::EraseTokensFromDB() {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  web_data_service_->RemoveAllTokens();
+}
+
+const bool TokenService::AreCredentialsValid() const {
   return !credentials_.lsid.empty() && !credentials_.sid.empty();
 }
 
-bool TokenService::HasLsid() const {
+const bool TokenService::HasLsid() const {
   return !credentials_.lsid.empty();
 }
 
@@ -33,18 +93,18 @@ const std::string& TokenService::GetLsid() const {
 }
 
 void TokenService::StartFetchingTokens() {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
   DCHECK(AreCredentialsValid());
-  sync_token_fetcher_->StartIssueAuthToken(credentials_.sid,
-                                           credentials_.lsid,
-                                           GaiaConstants::kSyncService);
-  talk_token_fetcher_->StartIssueAuthToken(credentials_.sid,
-                                           credentials_.lsid,
-                                           GaiaConstants::kTalkService);
+  for (int i = 0; i < kNumServices; i++) {
+    fetchers_[i]->StartIssueAuthToken(credentials_.sid,
+                                      credentials_.lsid,
+                                      kServices[i]);
+  }
 }
 
 // Services dependent on a token will check if a token is available.
 // If it isn't, they'll go to sleep until they get a token event.
-bool TokenService::HasTokenForService(const char* const service) const {
+const bool TokenService::HasTokenForService(const char* const service) const {
   return token_map_.count(service) > 0;
 }
 
@@ -52,28 +112,96 @@ const std::string& TokenService::GetTokenForService(
     const char* const service) const {
 
   if (token_map_.count(service) > 0) {
-    // map[key] is not const
+    // Note map[key] is not const.
     return (*token_map_.find(service)).second;
   }
   return EmptyString();
 }
 
-void TokenService::OnIssueAuthTokenSuccess(const std::string& service,
-                                           const std::string& auth_token) {
-  LOG(INFO) << "Got an authorization token for " << service;
-  token_map_[service] = auth_token;
+// Note that this can fire twice or more for any given service.
+// It can fire once from the DB read, and then once from the initial
+// fetcher. Future fetches can cause more notification firings.
+// The DB read will not however fire a notification if the fetcher
+// returned first. So it's always safe to use the latest notification.
+void TokenService::FireTokenAvailableNotification(
+    const std::string& service,
+    const std::string& auth_token) {
+
   TokenAvailableDetails details(service, auth_token);
   NotificationService::current()->Notify(
       NotificationType::TOKEN_AVAILABLE,
       Source<TokenService>(this),
       Details<const TokenAvailableDetails>(&details));
 }
-void TokenService::OnIssueAuthTokenFailure(const std::string& service,
-                                           const GaiaAuthError& error) {
-  LOG(WARNING) << "Auth token issuing failed for service:" << service;
+
+void TokenService::FireTokenRequestFailedNotification(
+    const std::string& service,
+    const GaiaAuthError& error) {
+
   TokenRequestFailedDetails details(service, error);
   NotificationService::current()->Notify(
       NotificationType::TOKEN_REQUEST_FAILED,
       Source<TokenService>(this),
       Details<const TokenRequestFailedDetails>(&details));
 }
+
+void TokenService::OnIssueAuthTokenSuccess(const std::string& service,
+                                           const std::string& auth_token) {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  LOG(INFO) << "Got an authorization token for " << service;
+  token_map_[service] = auth_token;
+  FireTokenAvailableNotification(service, auth_token);
+  SaveAuthTokenToDB(service, auth_token);
+}
+
+void TokenService::OnIssueAuthTokenFailure(const std::string& service,
+                                           const GaiaAuthError& error) {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  LOG(WARNING) << "Auth token issuing failed for service:" << service;
+  FireTokenRequestFailedNotification(service, error);
+}
+
+void TokenService::OnWebDataServiceRequestDone(WebDataService::Handle h,
+                                               const WDTypedResult* result) {
+  DCHECK(ChromeThread::CurrentlyOn(ChromeThread::UI));
+  DCHECK(token_loading_query_);
+  token_loading_query_ = NULL;
+
+  // If the fetch failed, there will be no result. In that case, we just don't
+  // load any tokens at all from the DB.
+  if (result) {
+    DCHECK(result->GetType() == TOKEN_RESULT);
+    const WDResult<std::map<std::string, std::string> > * token_result =
+        static_cast<const WDResult<std::map<std::string, std::string> > * > (
+            result);
+    LoadTokensIntoMemory(token_result->GetValue(), &token_map_);
+  }
+}
+
+// Load tokens from the db_token map into the in memory token map.
+void TokenService::LoadTokensIntoMemory(
+    const std::map<std::string, std::string>& db_tokens,
+    std::map<std::string, std::string>* in_memory_tokens) {
+
+  for (int i = 0; i < kNumServices; i++) {
+    // OnIssueAuthTokenSuccess should come from the same thread.
+    // If a token is already present in the map, it could only have
+    // come from a DB read or from IssueAuthToken. Since we should never
+    // fetch from the DB twice in a browser session, it must be from
+    // OnIssueAuthTokenSuccess, which is a live fetcher.
+    //
+    // Network fetched tokens take priority over DB tokens, so exclude tokens
+    // which have already been loaded by the fetcher.
+    if (!in_memory_tokens->count(kServices[i]) &&
+        db_tokens.count(kServices[i])) {
+      std::string db_token = db_tokens.find(kServices[i])->second;
+      if (!db_token.empty()) {
+        LOG(INFO) << "Loading " << kServices[i] << "token from DB:"
+            << db_token;
+        (*in_memory_tokens)[kServices[i]] = db_token;
+        FireTokenAvailableNotification(kServices[i], db_token);
+        // Failures are only for network errors.
+      }
+    }
+  }
+}
diff --git a/chrome/browser/net/gaia/token_service.h b/chrome/browser/net/gaia/token_service.h
index cf13699..42cb0ea 100644
--- a/chrome/browser/net/gaia/token_service.h
+++ b/chrome/browser/net/gaia/token_service.h
@@ -3,7 +3,33 @@
 // found in the LICENSE file.
 //
 // The TokenService will supply authentication tokens for any service that
-// needs it, such as sync.
+// needs it, such as sync. Whenever the user logs in, a controller watching
+// the token service is expected to call ClientLogin to derive a new SID and
+// LSID. Whenever such credentials are available, the TokenService should be
+// updated with new credentials. The controller should then start fetching
+// tokens, which will be written to the database after retrieval, as well as
+// provided to listeners.
+//
+// A token service controller like the ChromiumOS login is expected to:
+//
+// Initialize()  // Soon as you can
+// LoadTokensFromDB()  // When it's OK to talk to the database
+// UpdateCredentials()  // When user logs in
+// StartFetchingTokens()  // When it's safe to start fetching
+//
+// Typically a user of the TokenService is expected just to call:
+//
+// if (token_service.HasTokenForService(servicename)) {
+//   SetMyToken(token_service.GetTokenForService(servicename));
+// }
+// RegisterSomeObserver(token_service);
+//
+// Whenever a token update occurs:
+// OnTokenAvailable(...) {
+//   if (IsServiceICareAbout(notification.service())) {
+//     SetMyToken(notification.token())
+//   }
+// }
 
 #ifndef CHROME_BROWSER_NET_GAIA_TOKEN_SERVICE_H_
 #define CHROME_BROWSER_NET_GAIA_TOKEN_SERVICE_H_
@@ -11,18 +37,22 @@
 
 #include <map>
 #include <string>
+#include <vector>
 #include "base/scoped_ptr.h"
 #include "chrome/common/net/gaia/gaia_auth_consumer.h"
 #include "chrome/common/net/gaia/gaia_authenticator2.h"
 #include "chrome/test/testing_profile.h"
+#include "base/gtest_prod_util.h"
 
 class URLRequestContextGetter;
 
 // The TokenService is a Profile member, so all calls are expected
 // from the UI thread.
-class TokenService : public GaiaAuthConsumer {
+class TokenService : public GaiaAuthConsumer,
+                     public WebDataServiceConsumer {
  public:
-   TokenService() {}
+   TokenService();
+   virtual ~TokenService();
 
   // Notification classes
   class TokenAvailableDetails {
@@ -52,45 +82,94 @@ class TokenService : public GaiaAuthConsumer {
   };
 
   // Initialize this token service with a request source
-  // (usually from a GaiaAuthConsumer constant), a getter and
-  // results from ClientLogin.
-  void Initialize(
-      const char* const source,
-      URLRequestContextGetter* getter,
+  // (usually from a GaiaAuthConsumer constant), and the profile.
+  // Typically you'd then update the credentials.
+  void Initialize(const char* const source, Profile* profile);
+
+  // Update the credentials in the token service.
+  // Afterwards you can StartFetchingTokens.
+  void UpdateCredentials(
       const GaiaAuthConsumer::ClientLoginResult& credentials);
 
+  // Terminate any running requests and reset the TokenService to a clean
+  // slate. Resets in memory structures. Does not modify the DB.
+  // When this is done, no tokens will be left in memory and no
+  // user credentials will be left. Useful if a user is logging out.
+  // Initialize doesn't need to be called again but UpdateCredentials does.
+  void ResetCredentialsInMemory();
+
+  // Async load all tokens for services we know of from the DB.
+  // You should do this at startup. Optionally you can do it again
+  // after you reset in memory credentials.
+  void LoadTokensFromDB();
+
+  // Clear all DB stored tokens for the current profile. Tokens may still be
+  // available in memory. If a DB load is pending it may still be serviced.
+  void EraseTokensFromDB();
+
   // For legacy services with their own auth routines, they can just read
   // the LSID out directly. Deprecated.
-  bool HasLsid() const;
+  const bool HasLsid() const;
   const std::string& GetLsid() const;
+  // Did we get a proper LSID?
+  const bool AreCredentialsValid() const;
 
-  // On login, StartFetchingTokens() should be called to kick off token
-  // fetching.
   // Tokens will be fetched for all services(sync, talk) in the background.
-  // Results come back via event channel. Services can also poll.
+  // Results come back via event channel. Services can also poll before events
+  // are issued.
   void StartFetchingTokens();
-  bool HasTokenForService(const char* const service) const;
+  const bool HasTokenForService(const char* const service) const;
   const std::string& GetTokenForService(const char* const service) const;
 
-  // Callbacks from the fetchers.
-  void OnIssueAuthTokenSuccess(const std::string& service,
-                               const std::string& auth_token);
-  void OnIssueAuthTokenFailure(const std::string& service,
-                               const GaiaAuthError& error);
+  // GaiaAuthConsumer implementation.
+  virtual void OnIssueAuthTokenSuccess(const std::string& service,
+                                       const std::string& auth_token);
+  virtual void OnIssueAuthTokenFailure(const std::string& service,
+                                       const GaiaAuthError& error);
+
+  // WebDataServiceConsumer implementation.
+  virtual void OnWebDataServiceRequestDone(WebDataService::Handle h,
+                                           const WDTypedResult* result);
 
  private:
-  // Did we get a proper LSID?
-  bool AreCredentialsValid() const;
+
+  void FireTokenAvailableNotification(const std::string& service,
+                                      const std::string& auth_token);
+
+  void FireTokenRequestFailedNotification(const std::string& service,
+                                          const GaiaAuthError& error);
+
+  void LoadTokensIntoMemory(const std::map<std::string, std::string>& in_toks,
+                            std::map<std::string, std::string>* out_toks);
+
+  void SaveAuthTokenToDB(const std::string& service,
+                         const std::string& auth_token);
+
+  // Web data service to access tokens from.
+  scoped_refptr<WebDataService> web_data_service_;
+  // Getter to use for fetchers.
+  scoped_refptr<URLRequestContextGetter> getter_;
+  // Request handle to load Gaia tokens from DB.
+  WebDataService::Handle token_loading_query_;
 
   // Gaia request source for Gaia accounting.
   std::string source_;
   // Credentials from ClientLogin for Issuing auth tokens.
   GaiaAuthConsumer::ClientLoginResult credentials_;
-  scoped_ptr<GaiaAuthenticator2> sync_token_fetcher_;
-  scoped_ptr<GaiaAuthenticator2> talk_token_fetcher_;
+
+  // Size of array of services (must be defined here).
+  static const int kNumServices = 2;
+  // List of services that we're performing operations for.
+  static const char* kServices[kNumServices];
+  // A bunch of fetchers suitable for token issuing. We don't care about
+  // the ordering, nor do we care which is for which service.
+  scoped_ptr<GaiaAuthenticator2> fetchers_[kNumServices];
   // Map from service to token.
   std::map<std::string, std::string> token_map_;
 
+  FRIEND_TEST_ALL_PREFIXES(TokenServiceTest, LoadTokensIntoMemoryBasic);
+  FRIEND_TEST_ALL_PREFIXES(TokenServiceTest, LoadTokensIntoMemoryAdvanced);
+
   DISALLOW_COPY_AND_ASSIGN(TokenService);
 };
 
diff --git a/chrome/browser/net/gaia/token_service_unittest.cc b/chrome/browser/net/gaia/token_service_unittest.cc
index 0bd7e56..a9c6863 100644
--- a/chrome/browser/net/gaia/token_service_unittest.cc
+++ b/chrome/browser/net/gaia/token_service_unittest.cc
@@ -4,12 +4,17 @@
 //
 // This file defines a unit test for the profile's token service.
 
+#include "base/scoped_temp_dir.h"
 #include "chrome/browser/net/gaia/token_service.h"
-#include "chrome/common/net/gaia/gaia_authenticator2_unittest.h"
+#include "chrome/browser/password_manager/encryptor.h"
+#include "chrome/browser/webdata/web_data_service.h"
 #include "chrome/common/net/gaia/gaia_auth_consumer.h"
+#include "chrome/common/net/gaia/gaia_authenticator2_unittest.h"
 #include "chrome/common/net/gaia/gaia_constants.h"
-#include "testing/gtest/include/gtest/gtest.h"
+#include "chrome/common/net/test_url_fetcher_factory.h"
+#include "chrome/test/signaling_task.h"
 #include "chrome/test/test_notification_tracker.h"
+#include "testing/gtest/include/gtest/gtest.h"
 
 // TestNotificationTracker doesn't do a deep copy on the notification details.
 // We have to in order to read it out, or we have a bad ptr, since the details
@@ -56,27 +61,73 @@ class TokenFailedTracker : public TestNotificationTracker {
 
 class TokenServiceTest : public testing::Test {
  public:
-  TokenServiceTest() {
+  TokenServiceTest()
+      : ui_thread_(ChromeThread::UI, &message_loop_),
+        db_thread_(ChromeThread::DB) {
+  }
+
+  virtual void SetUp() {
+#if defined(OS_MACOSX)
+    Encryptor::UseMockKeychain(true);
+#endif
     credentials_.sid = "sid";
     credentials_.lsid = "lsid";
     credentials_.token = "token";
     credentials_.data = "data";
 
+    ASSERT_TRUE(temp_.CreateUniqueTempDir());
+    ASSERT_TRUE(db_thread_.Start());
+
+    // Testing profile responsible for deleting the temp dir.
+    profile_.reset(new TestingProfile(temp_.Take()));
+    profile_->CreateWebDataService(false);
+    WaitForDBLoadCompletion();
+
     success_tracker_.ListenFor(NotificationType::TOKEN_AVAILABLE,
                                Source<TokenService>(&service_));
     failure_tracker_.ListenFor(NotificationType::TOKEN_REQUEST_FAILED,
                                Source<TokenService>(&service_));
 
-    service_.Initialize("test", profile_.GetRequestContext(), credentials_);
+    service_.Initialize("test", profile_.get());
+    service_.UpdateCredentials(credentials_);
+
+    URLFetcher::set_factory(NULL);
+  }
+
+  virtual void TearDown() {
+    // You have to destroy the profile before the db_thread_ stops.
+    if (profile_.get()) {
+      profile_.reset(NULL);
+    }
+
+    db_thread_.Stop();
+    MessageLoop::current()->PostTask(FROM_HERE, new MessageLoop::QuitTask);
+    MessageLoop::current()->Run();
   }
 
+  void WaitForDBLoadCompletion() {
+    // The WebDB does all work on the DB thread. This will add an event
+    // to the end of the DB thread, so when we reach this task, all DB
+    // operations should be complete.
+    WaitableEvent done(false, false);
+    ChromeThread::PostTask(
+        ChromeThread::DB, FROM_HERE, new SignalingTask(&done));
+    done.Wait();
+
+    // Notifications should be returned from the DB thread onto the UI thread.
+    message_loop_.RunAllPending();
+  }
+
+  MessageLoopForUI message_loop_;
+  ChromeThread ui_thread_;  // Mostly so DCHECKS pass.
+  ChromeThread db_thread_;  // WDS on here
+
   TokenService service_;
   TokenAvailableTracker success_tracker_;
   TokenFailedTracker failure_tracker_;
   GaiaAuthConsumer::ClientLoginResult credentials_;
-
- private:
-  TestingProfile profile_;
+  scoped_ptr<TestingProfile> profile_;
+  ScopedTempDir temp_;
 };
 
 TEST_F(TokenServiceTest, SanityCheck) {
@@ -99,7 +150,7 @@ TEST_F(TokenServiceTest, NotificationSuccess) {
 
   TokenService::TokenAvailableDetails details =
       success_tracker_.get_last_token_details();
-  // MSVC doesn't like this comparison as EQ
+  // MSVC doesn't like this comparison as EQ.
   EXPECT_TRUE(details.service() == GaiaConstants::kSyncService);
   EXPECT_EQ(details.token(), "token");
 }
@@ -116,9 +167,9 @@ TEST_F(TokenServiceTest, NotificationFailed) {
   TokenService::TokenRequestFailedDetails details =
       failure_tracker_.get_last_token_details();
 
-  // MSVC doesn't like this comparison as EQ
+  // MSVC doesn't like this comparison as EQ.
   EXPECT_TRUE(details.service() == GaiaConstants::kSyncService);
-  EXPECT_TRUE(details.error() == error);  // Struct has no print function
+  EXPECT_TRUE(details.error() == error);  // Struct has no print function.
 }
 
 TEST_F(TokenServiceTest, OnTokenSuccessUpdate) {
@@ -144,15 +195,54 @@ TEST_F(TokenServiceTest, OnTokenSuccess) {
   // result with ClientLogin, it doesn't matter, we should still get it back.
   EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kSyncService), "token");
 
-  // Check 2nd service
+  // Check the second service.
   service_.OnIssueAuthTokenSuccess(GaiaConstants::kTalkService, "token2");
   EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kTalkService));
   EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kTalkService), "token2");
 
-  // Didn't change
+  // It didn't change.
   EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kSyncService), "token");
 }
 
+TEST_F(TokenServiceTest, ResetSimple) {
+  service_.OnIssueAuthTokenSuccess(GaiaConstants::kSyncService, "token");
+  EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_TRUE(service_.HasLsid());
+
+  service_.ResetCredentialsInMemory();
+
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_FALSE(service_.HasLsid());
+}
+
+TEST_F(TokenServiceTest, ResetComplex) {
+  TestURLFetcherFactory factory;
+  URLFetcher::set_factory(&factory);
+  service_.StartFetchingTokens();
+  // You have to call delegates by hand with the test fetcher,
+  // Let's pretend only one returned.
+
+  service_.OnIssueAuthTokenSuccess(GaiaConstants::kSyncService, "eraseme");
+  EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kSyncService),
+            "eraseme");
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kTalkService));
+
+  service_.ResetCredentialsInMemory();
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kTalkService));
+
+  // Now start using it again.
+  service_.UpdateCredentials(credentials_);
+  service_.StartFetchingTokens();
+
+  service_.OnIssueAuthTokenSuccess(GaiaConstants::kSyncService, "token");
+  service_.OnIssueAuthTokenSuccess(GaiaConstants::kTalkService, "token2");
+
+  EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kSyncService), "token");
+  EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kTalkService), "token2");
+}
+
 TEST_F(TokenServiceTest, FullIntegration) {
   MockFactory<MockFetcher> factory;
   std::string result = "SID=sid\nLSID=lsid\nAuth=auth\n";
@@ -169,4 +259,116 @@ TEST_F(TokenServiceTest, FullIntegration) {
   // result with ClientLogin, it doesn't matter, we should still get it back.
   EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kSyncService), result);
   EXPECT_EQ(service_.GetTokenForService(GaiaConstants::kTalkService), result);
+
+  service_.ResetCredentialsInMemory();
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kTalkService));
+}
+
+TEST_F(TokenServiceTest, LoadTokensIntoMemoryBasic) {
+  // Validate that the method sets proper data in notifications and map.
+  std::map<std::string, std::string> db_tokens;
+  std::map<std::string, std::string> memory_tokens;
+
+  service_.LoadTokensIntoMemory(db_tokens, &memory_tokens);
+  EXPECT_TRUE(db_tokens.empty());
+  EXPECT_TRUE(memory_tokens.empty());
+  EXPECT_EQ(0U, success_tracker_.size());
+
+  db_tokens[GaiaConstants::kSyncService] = "token";
+  service_.LoadTokensIntoMemory(db_tokens, &memory_tokens);
+  EXPECT_EQ(1U, success_tracker_.size());
+
+  TokenService::TokenAvailableDetails details =
+      success_tracker_.get_last_token_details();
+  // MSVC doesn't like this comparison as EQ.
+  EXPECT_TRUE(details.service() == GaiaConstants::kSyncService);
+  EXPECT_EQ(details.token(), "token");
+  EXPECT_EQ(1U, memory_tokens.count(GaiaConstants::kSyncService));
+  EXPECT_EQ(memory_tokens[GaiaConstants::kSyncService], "token");
+}
+
+TEST_F(TokenServiceTest, LoadTokensIntoMemoryAdvanced) {
+  // LoadTokensIntoMemory should avoid setting tokens already in the
+  // token map.
+  std::map<std::string, std::string> db_tokens;
+  std::map<std::string, std::string> memory_tokens;
+
+  db_tokens["ignore"] = "token";
+
+  service_.LoadTokensIntoMemory(db_tokens, &memory_tokens);
+  EXPECT_TRUE(memory_tokens.empty());
+  db_tokens[GaiaConstants::kSyncService] = "pepper";
+
+  service_.LoadTokensIntoMemory(db_tokens, &memory_tokens);
+  EXPECT_EQ(1U, memory_tokens.count(GaiaConstants::kSyncService));
+  EXPECT_EQ(memory_tokens[GaiaConstants::kSyncService], "pepper");
+  EXPECT_EQ(1U, success_tracker_.size());
+  success_tracker_.Reset();
+
+  // SyncService token is already in memory. Pretend we got it off
+  // the disk as well, but an older token.
+  db_tokens[GaiaConstants::kSyncService] = "ignoreme";
+  db_tokens[GaiaConstants::kTalkService] = "tomato";
+  service_.LoadTokensIntoMemory(db_tokens, &memory_tokens);
+
+  EXPECT_EQ(2U, memory_tokens.size());
+  EXPECT_EQ(1U, memory_tokens.count(GaiaConstants::kTalkService));
+  EXPECT_EQ(memory_tokens[GaiaConstants::kTalkService], "tomato");
+  EXPECT_EQ(1U, success_tracker_.size());
+  EXPECT_EQ(1U, memory_tokens.count(GaiaConstants::kSyncService));
+  EXPECT_EQ(memory_tokens[GaiaConstants::kSyncService], "pepper");
+}
+
+TEST_F(TokenServiceTest, WebDBLoadIntegration) {
+  service_.LoadTokensFromDB();
+  WaitForDBLoadCompletion();
+  EXPECT_EQ(0U, success_tracker_.size());
+
+  // Should result in DB write.
+  service_.OnIssueAuthTokenSuccess(GaiaConstants::kSyncService, "token");
+  EXPECT_EQ(1U, success_tracker_.size());
+
+  EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  // Clean slate.
+  service_.ResetCredentialsInMemory();
+  success_tracker_.Reset();
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kSyncService));
+
+  service_.LoadTokensFromDB();
+  WaitForDBLoadCompletion();
+
+  EXPECT_EQ(1U, success_tracker_.size());
+  EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kTalkService));
+  EXPECT_FALSE(service_.HasLsid());
+}
+
+TEST_F(TokenServiceTest, MultipleLoadResetIntegration) {
+  // Should result in DB write.
+  service_.OnIssueAuthTokenSuccess(GaiaConstants::kSyncService, "token");
+  service_.ResetCredentialsInMemory();
+  success_tracker_.Reset();
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kSyncService));
+
+  service_.LoadTokensFromDB();
+  WaitForDBLoadCompletion();
+
+  service_.LoadTokensFromDB();  // Should do nothing.
+  WaitForDBLoadCompletion();
+
+  EXPECT_EQ(1U, success_tracker_.size());
+  EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kSyncService));
+  EXPECT_FALSE(service_.HasTokenForService(GaiaConstants::kTalkService));
+  EXPECT_FALSE(service_.HasLsid());
+
+  // Reset it one more time so there's no surprises.
+  service_.ResetCredentialsInMemory();
+  success_tracker_.Reset();
+
+  service_.LoadTokensFromDB();
+  WaitForDBLoadCompletion();
+
+  EXPECT_EQ(1U, success_tracker_.size());
+  EXPECT_TRUE(service_.HasTokenForService(GaiaConstants::kSyncService));
 }