Identity API: add multi-account support to token cache and request queues

chrome/browser/extensions/api/identity/identity_apitest.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/command_line.h"
 #include "base/strings/string_util.h"
 #include "base/strings/stringprintf.h"
 #include "base/values.h"
 #include "chrome/browser/chrome_notification_types.h"
 #include "chrome/browser/extensions/api/identity/identity_api.h"
 #include "chrome/browser/extensions/component_loader.h"
 #include "chrome/browser/extensions/extension_apitest.h"
 #include "chrome/browser/extensions/extension_browsertest.h"
 #include "chrome/browser/extensions/extension_function_test_utils.h"
 #include "chrome/browser/extensions/extension_service.h"
 #include "chrome/browser/profiles/profile.h"
+#include "chrome/browser/signin/profile_oauth2_token_service.h"
+#include "chrome/browser/signin/profile_oauth2_token_service_factory.h"
 #include "chrome/browser/ui/browser.h"
 #include "chrome/browser/ui/browser_window.h"
 #include "chrome/common/chrome_switches.h"
 #include "chrome/common/extensions/api/identity/oauth2_manifest_handler.h"
 #include "chrome/test/base/in_process_browser_test.h"
 #include "chrome/test/base/test_switches.h"
 #include "content/public/browser/notification_service.h"
 #include "content/public/browser/notification_source.h"
 #include "content/public/test/test_utils.h"
 #include "extensions/common/id_util.h"
@@ skipping 322 matching lines @@
     else
       ext = LoadExtensionAsComponent(component_manifest_path);
     OAuth2Info& oauth2_info =
         const_cast<OAuth2Info&>(OAuth2Info::GetOAuth2Info(ext));
     if ((fields_to_set & CLIENT_ID) != 0)
       oauth2_info.client_id = "client1";
     if ((fields_to_set & SCOPES) != 0) {
       oauth2_info.scopes.push_back("scope1");
       oauth2_info.scopes.push_back("scope2");
     }
+
+    extension_id_ = ext->id();
+    oauth_scopes_ = std::set<std::string>(oauth2_info.scopes.begin(),
+                                          oauth2_info.scopes.end());
     return ext;
   }
 
   IdentityAPI* id_api() {
     return IdentityAPI::GetFactoryInstance()->GetForProfile(
         browser()->profile());
   }
+
+  const std::string GetPrimaryAccountId() {
+    ProfileOAuth2TokenService* token_service =
+        ProfileOAuth2TokenServiceFactory::GetForProfile(browser()->profile());
+    return token_service->GetPrimaryAccountId();
+  }
+
+  void SetCachedToken(const IdentityTokenCacheValue& token_data) {
+    ExtensionTokenKey key(extension_id_, GetPrimaryAccountId(), oauth_scopes_);
+    id_api()->SetCachedToken(key, token_data);
+  }
+
+  const IdentityTokenCacheValue& GetCachedToken() {
+    ExtensionTokenKey key(extension_id_, GetPrimaryAccountId(), oauth_scopes_);
+    return id_api()->GetCachedToken(key);
+  }
+
+  void QueueRequestStart(IdentityMintRequestQueue::MintType type,
+                         IdentityMintRequestQueue::Request* request) {
+    ExtensionTokenKey key(extension_id_, GetPrimaryAccountId(), oauth_scopes_);
+    id_api()->mint_queue()->RequestStart(type, key, request);
+  }
+
+  void QueueRequestComplete(IdentityMintRequestQueue::MintType type,
+                            IdentityMintRequestQueue::Request* request) {
+    ExtensionTokenKey key(extension_id_, GetPrimaryAccountId(), oauth_scopes_);
+    id_api()->mint_queue()->RequestComplete(type, key, request);
+  }
+
+ private:
+  std::string extension_id_;
+  std::set<std::string> oauth_scopes_;
 };
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        NoClientId) {
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(CreateExtension(SCOPES));
   std::string error = utils::RunFunctionAndReturnError(
       func.get(), "[{}]", browser());
   EXPECT_EQ(std::string(errors::kInvalidClientId), error);
   EXPECT_FALSE(func->login_ui_shown());
@@ skipping 59 matching lines @@
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::ISSUE_ADVICE_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_)).WillOnce(Return(flow));
   std::string error = utils::RunFunctionAndReturnError(
       func.get(), "[{}]", browser());
   EXPECT_EQ(std::string(errors::kNoGrant), error);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
-  EXPECT_EQ(
-      IdentityTokenCacheValue::CACHE_STATUS_ADVICE,
-      id_api()->GetCachedToken(extension->id(), oauth2_info.scopes).status());
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_ADVICE,
+            GetCachedToken().status());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        NonInteractiveMintBadCredentials) {
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(CreateExtension(CLIENT_ID | SCOPES));
   EXPECT_CALL(*func.get(), HasLoginToken())
       .WillOnce(Return(true));
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::MINT_TOKEN_BAD_CREDENTIALS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_)).WillOnce(Return(flow));
   std::string error = utils::RunFunctionAndReturnError(
       func.get(), "[{}]", browser());
   EXPECT_TRUE(StartsWithASCII(error, errors::kAuthFailure, false));
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        NonInteractiveSuccess) {
 #if defined(OS_WIN) && defined(USE_ASH)
   // Disable this test in Metro+Ash for now (http://crbug.com/262796).
   if (CommandLine::ForCurrentProcess()->HasSwitch(switches::kAshBrowserTests))
     return;
 #endif
 
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   func->set_extension(extension.get());
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::MINT_TOKEN_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_)).WillOnce(Return(flow));
   scoped_ptr<base::Value> value(utils::RunFunctionAndReturnSingleResult(
       func.get(), "[{}]", browser()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
   EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
-            id_api()->GetCachedToken(extension->id(),
-                                     oauth2_info.scopes).status());
+            GetCachedToken().status());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        InteractiveLoginCanceled) {
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(CreateExtension(CLIENT_ID | SCOPES));
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(false));
   func->set_login_ui_result(false);
   std::string error = utils::RunFunctionAndReturnError(
       func.get(), "[{\"interactive\": true}]", browser());
@@ skipping 218 matching lines @@
         func.get(), "[{\"interactive\": true}]", browser());
     EXPECT_EQ(it->second, error);
     EXPECT_FALSE(func->login_ui_shown());
     EXPECT_TRUE(func->scope_ui_shown());
   }
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        InteractiveApprovalSuccess) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::ISSUE_ADVICE_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_))
       .WillOnce(Return(flow));
 
   scoped_ptr<base::Value> value(utils::RunFunctionAndReturnSingleResult(
       func.get(), "[{\"interactive\": true}]", browser()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_TRUE(func->scope_ui_shown());
 
   EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
-            id_api()->GetCachedToken(extension->id(),
-                                     oauth2_info.scopes).status());
+            GetCachedToken().status());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest, NoninteractiveQueue) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
 
   // Create a fake request to block the queue.
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
-  std::set<std::string> scopes(oauth2_info.scopes.begin(),
-                               oauth2_info.scopes.end());
-  IdentityAPI* id_api =
-      extensions::IdentityAPI::GetFactoryInstance()->GetForProfile(
-          browser()->profile());
-  IdentityMintRequestQueue* queue = id_api->mint_queue();
   MockQueuedMintRequest queued_request;
   IdentityMintRequestQueue::MintType type =
       IdentityMintRequestQueue::MINT_TYPE_NONINTERACTIVE;
 
   EXPECT_CALL(queued_request, StartMintToken(type)).Times(1);
-  queue->RequestStart(type, extension->id(), scopes, &queued_request);
+  QueueRequestStart(type, &queued_request);
 
   // The real request will start processing, but wait in the queue behind
   // the blocker.
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   RunFunctionAsync(func.get(), "[{}]");
   // Verify that we have fetched the login token at this point.
   testing::Mock::VerifyAndClearExpectations(func.get());
 
   // The flow will be created after the first queued request clears.
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::MINT_TOKEN_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_)).WillOnce(Return(flow));
 
-  queue->RequestComplete(type, extension->id(), scopes, &queued_request);
+  QueueRequestComplete(type, &queued_request);
 
   scoped_ptr<base::Value> value(WaitForSingleResult(func.get()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest, InteractiveQueue) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
 
   // Create a fake request to block the queue.
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
-  std::set<std::string> scopes(oauth2_info.scopes.begin(),
-                               oauth2_info.scopes.end());
-  IdentityAPI* id_api =
-      extensions::IdentityAPI::GetFactoryInstance()->GetForProfile(
-          browser()->profile());
-  IdentityMintRequestQueue* queue = id_api->mint_queue();
   MockQueuedMintRequest queued_request;
   IdentityMintRequestQueue::MintType type =
       IdentityMintRequestQueue::MINT_TYPE_INTERACTIVE;
 
   EXPECT_CALL(queued_request, StartMintToken(type)).Times(1);
-  queue->RequestStart(type, extension->id(), scopes, &queued_request);
+  QueueRequestStart(type, &queued_request);
 
   // The real request will start processing, but wait in the queue behind
   // the blocker.
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   TestOAuth2MintTokenFlow* flow1 = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::ISSUE_ADVICE_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_)).WillOnce(Return(flow1));
   RunFunctionAsync(func.get(), "[{\"interactive\": true}]");
   // Verify that we have fetched the login token and run the first flow.
   testing::Mock::VerifyAndClearExpectations(func.get());
   EXPECT_FALSE(func->scope_ui_shown());
 
   // The UI will be displayed and a token retrieved after the first
   // queued request clears.
-  queue->RequestComplete(type, extension->id(), scopes, &queued_request);
+  QueueRequestComplete(type, &queued_request);
 
   scoped_ptr<base::Value> value(WaitForSingleResult(func.get()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_TRUE(func->scope_ui_shown());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        InteractiveQueuedNoninteractiveFails) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
 
   // Create a fake request to block the interactive queue.
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
-  std::set<std::string> scopes(oauth2_info.scopes.begin(),
-                               oauth2_info.scopes.end());
-  IdentityAPI* id_api =
-      extensions::IdentityAPI::GetFactoryInstance()->GetForProfile(
-          browser()->profile());
-  IdentityMintRequestQueue* queue = id_api->mint_queue();
   MockQueuedMintRequest queued_request;
   IdentityMintRequestQueue::MintType type =
       IdentityMintRequestQueue::MINT_TYPE_INTERACTIVE;
 
   EXPECT_CALL(queued_request, StartMintToken(type)).Times(1);
-  queue->RequestStart(type, extension->id(), scopes, &queued_request);
+  QueueRequestStart(type, &queued_request);
 
   // Non-interactive requests fail without hitting GAIA, because a
   // consent UI is known to be up.
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   std::string error = utils::RunFunctionAndReturnError(
       func.get(), "[{}]", browser());
   EXPECT_EQ(std::string(errors::kNoGrant), error);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 
-  queue->RequestComplete(type, extension->id(), scopes, &queued_request);
+  QueueRequestComplete(type, &queued_request);
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        NonInteractiveCacheHit) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
 
   // pre-populate the cache with a token
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
   IdentityTokenCacheValue token(kAccessToken,
                                 base::TimeDelta::FromSeconds(3600));
-  id_api()->SetCachedToken(extension->id(), oauth2_info.scopes, token);
+  SetCachedToken(token);
 
   // Get a token. Should not require a GAIA request.
   EXPECT_CALL(*func.get(), HasLoginToken())
       .WillOnce(Return(true));
   scoped_ptr<base::Value> value(utils::RunFunctionAndReturnSingleResult(
       func.get(), "[{}]", browser()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        NonInteractiveIssueAdviceCacheHit) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
 
   // pre-populate the cache with advice
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
   IssueAdviceInfo info;
   IdentityTokenCacheValue token(info);
-  id_api()->SetCachedToken(extension->id(), oauth2_info.scopes, token);
+  SetCachedToken(token);
 
   // Should return an error without a GAIA request.
   EXPECT_CALL(*func.get(), HasLoginToken())
       .WillOnce(Return(true));
   std::string error = utils::RunFunctionAndReturnError(
       func.get(), "[{}]", browser());
   EXPECT_EQ(std::string(errors::kNoGrant), error);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        InteractiveCacheHit) {
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   func->set_extension(extension.get());
 
   // Create a fake request to block the queue.
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
-  std::set<std::string> scopes(oauth2_info.scopes.begin(),
-                               oauth2_info.scopes.end());
-  IdentityMintRequestQueue* queue = id_api()->mint_queue();
   MockQueuedMintRequest queued_request;
   IdentityMintRequestQueue::MintType type =
       IdentityMintRequestQueue::MINT_TYPE_INTERACTIVE;
 
   EXPECT_CALL(queued_request, StartMintToken(type)).Times(1);
-  queue->RequestStart(type, extension->id(), scopes, &queued_request);
+  QueueRequestStart(type, &queued_request);
 
   // The real request will start processing, but wait in the queue behind
   // the blocker.
   EXPECT_CALL(*func.get(), HasLoginToken()).WillOnce(Return(true));
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::ISSUE_ADVICE_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_)).WillOnce(Return(flow));
   RunFunctionAsync(func.get(), "[{\"interactive\": true}]");
 
   // Populate the cache with a token while the request is blocked.
   IdentityTokenCacheValue token(kAccessToken,
                                 base::TimeDelta::FromSeconds(3600));
-  id_api()->SetCachedToken(extension->id(), oauth2_info.scopes, token);
+  SetCachedToken(token);
 
   // When we wake up the request, it returns the cached token without
   // displaying a UI, or hitting GAIA.
 
-  queue->RequestComplete(type, extension->id(), scopes, &queued_request);
+  QueueRequestComplete(type, &queued_request);
 
   scoped_ptr<base::Value> value(WaitForSingleResult(func.get()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_FALSE(func->login_ui_shown());
   EXPECT_FALSE(func->scope_ui_shown());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
                        LoginInvalidatesTokenCache) {
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   scoped_refptr<const Extension> extension(CreateExtension(CLIENT_ID | SCOPES));
   func->set_extension(extension.get());
 
   // pre-populate the cache with a token
-  const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
   IdentityTokenCacheValue token(kAccessToken,
                                 base::TimeDelta::FromSeconds(3600));
-  id_api()->SetCachedToken(extension->id(), oauth2_info.scopes, token);
+  SetCachedToken(token);
 
   // Because the user is not signed in, the token will be removed,
   // and we'll hit GAIA for new tokens.
   EXPECT_CALL(*func.get(), HasLoginToken())
       .WillOnce(Return(false));
   func->set_login_ui_result(true);
   TestOAuth2MintTokenFlow* flow = new TestOAuth2MintTokenFlow(
       TestOAuth2MintTokenFlow::ISSUE_ADVICE_SUCCESS, func.get());
   EXPECT_CALL(*func.get(), CreateMintTokenFlow(_))
       .WillOnce(Return(flow));
 
   scoped_ptr<base::Value> value(utils::RunFunctionAndReturnSingleResult(
       func.get(), "[{\"interactive\": true}]", browser()));
   std::string access_token;
   EXPECT_TRUE(value->GetAsString(&access_token));
   EXPECT_EQ(std::string(kAccessToken), access_token);
   EXPECT_TRUE(func->login_ui_shown());
   EXPECT_TRUE(func->scope_ui_shown());
   EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
-            id_api()->GetCachedToken(extension->id(),
-                                     oauth2_info.scopes).status());
+            GetCachedToken().status());
 }
 
 IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest, ComponentWithChromeClientId) {
   scoped_refptr<MockGetAuthTokenFunction> func(new MockGetAuthTokenFunction());
   scoped_refptr<const Extension> extension(
       CreateExtension(SCOPES | AS_COMPONENT));
   func->set_extension(extension.get());
   const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension.get());
   EXPECT_TRUE(oauth2_info.client_id.empty());
   EXPECT_FALSE(func->GetOAuth2ClientId().empty());
@@ skipping 20 matching lines @@
         browser(),
         extension_function_test_utils::NONE);
   }
 
   IdentityAPI* id_api() {
     return IdentityAPI::GetFactoryInstance()->GetForProfile(
         browser()->profile());
   }
 
   void SetCachedToken(IdentityTokenCacheValue& token_data) {
-    id_api()->SetCachedToken(extensions::id_util::GenerateId(kExtensionId),
-                             std::vector<std::string>(), token_data);
+    ExtensionTokenKey key(extensions::id_util::GenerateId(kExtensionId),
+                          "test@example.com",
+                          std::set<std::string>());
+    id_api()->SetCachedToken(key, token_data);
   }
 
   const IdentityTokenCacheValue& GetCachedToken() {
     return id_api()->GetCachedToken(
-        extensions::id_util::GenerateId(kExtensionId),
-        std::vector<std::string>());
+        ExtensionTokenKey(extensions::id_util::GenerateId(kExtensionId),
+                          "test@example.com",
+                          std::set<std::string>()));
   }
 };
 
 IN_PROC_BROWSER_TEST_F(RemoveCachedAuthTokenFunctionTest, NotFound) {
   EXPECT_TRUE(InvalidateDefaultToken());
   EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
             GetCachedToken().status());
 }
 
 IN_PROC_BROWSER_TEST_F(RemoveCachedAuthTokenFunctionTest, Advice) {
@@ skipping 12 matching lines @@
   EXPECT_TRUE(InvalidateDefaultToken());
   EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
             GetCachedToken().status());
   EXPECT_EQ("non_matching_token", GetCachedToken().token());
 }
 
 IN_PROC_BROWSER_TEST_F(RemoveCachedAuthTokenFunctionTest, MatchingToken) {
   IdentityTokenCacheValue token(kAccessToken,
                                 base::TimeDelta::FromSeconds(3600));
   SetCachedToken(token);
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
+            GetCachedToken().status());
   EXPECT_TRUE(InvalidateDefaultToken());
   EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
             GetCachedToken().status());
 }
 
 class LaunchWebAuthFlowFunctionTest : public AsyncExtensionBrowserTest {
  public:
   virtual void SetUpCommandLine(CommandLine* command_line) OVERRIDE {
     // Reduce performance test variance by disabling background networking.
     command_line->AppendSwitch(switches::kDisableBackgroundNetworking);
@@ skipping 153 matching lines @@
   EXPECT_EQ(std::string("https://abcdefghij.chromiumapp.org/callback#test"),
             url);
 }
 
 }  // namespace extensions
 
 // Tests the chrome.identity API implemented by custom JS bindings .
 IN_PROC_BROWSER_TEST_F(ExtensionApiTest, ChromeIdentityJsBindings) {
   ASSERT_TRUE(RunExtensionTest("identity/js_bindings")) << message_;
 }