Add tests for redirect responses from SafeBrowsingProtocolManager.

chrome/browser/safe_browsing/protocol_manager_unittest.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/logging.h"
 #include "base/single_thread_task_runner.h"
 #include "base/stringprintf.h"
 #include "base/thread_task_runner_handle.h"
 #include "base/time.h"
@@ skipping 235 matching lines @@
 
 class MockProtocolDelegate : public SafeBrowsingProtocolManagerDelegate {
  public:
   MockProtocolDelegate() {}
   virtual ~MockProtocolDelegate() {}
 
   MOCK_METHOD0(UpdateStarted, void());
   MOCK_METHOD1(UpdateFinished, void(bool));
   MOCK_METHOD0(ResetDatabase, void());
   MOCK_METHOD1(GetChunks, void(GetChunksCallback));
-  MOCK_METHOD2(AddChunks, void(const std::string&, SBChunkList*));
+  MOCK_METHOD3(AddChunks, void(const std::string&, SBChunkList*,
+                               AddChunksCallback));
   MOCK_METHOD1(DeleteChunks, void(std::vector<SBChunkDelete>*));
 };
 
 // ImmediateSingleThreadTaskRunner will ignore delayed times for tasks.
 // This is primarily used to run the timer tasks immediately, and prevent
 // the need for constructing a MessageLoop.
 class ImmediateSingleThreadTaskRunner : public base::SingleThreadTaskRunner {
  public:
   virtual bool RunsTasksOnCurrentThread() const OVERRIDE {
     return true;
@@ skipping 38 matching lines @@
 // |InvokeGetChunksCallback| is required because GMock's InvokeArgument action
 // expects to use operator(), and a Callback only provides Run().
 // TODO(cbentzel): Use ACTION or ACTION_TEMPLATE instead?
 void InvokeGetChunksCallback(
     const std::vector<SBListChunkRanges>& ranges,
     bool database_error,
     SafeBrowsingProtocolManagerDelegate::GetChunksCallback callback) {
   callback.Run(ranges, database_error);
 }
 
+// |HandleAddChunks| deletes the chunks and asynchronously invokes
+// |callback| since SafeBrowsingProtocolManager is not re-entrant at the time
+// this is called. This guarantee is part of the
+// SafeBrowsingProtocolManagerDelegate contract.
+void HandleAddChunks(
+    const std::string& unused_list,
+    SBChunkList* chunks,
+    SafeBrowsingProtocolManagerDelegate::AddChunksCallback callback) {
+  delete chunks;
+  scoped_refptr<base::SingleThreadTaskRunner> task_runner(
+      base::ThreadTaskRunnerHandle::Get());
+  if (!task_runner)
+    return;
+  task_runner->PostTask(FROM_HERE, callback);
+}
+
 }  // namespace
 
 // Tests that the Update protocol will be skipped if there are problems
 // accessing the database.
 TEST_F(SafeBrowsingProtocolManagerTest, ProblemAccessingDatabase) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     true)));
   EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());

[mattm, 2012/12/06 00:51:40]

Could inserting some EXPECT_FALSE(pm->IsUpdateScheduled()); in strategic places
also be useful?

[cbentzel, 2012/12/17 11:24:10]

I put it in a few locations (such as the multiple redirect responses). Not sure
if it helps that much.

 }
 
-// Tests the contents of the POST body when the local database is empty.
-TEST_F(SafeBrowsingProtocolManagerTest, EmptyDatabase) {
+// Tests what happens when there is a response with no chunks.
+TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseEmptyBody) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, but an empty body.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString("");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
 // Tests the contents of the POST body when there are contents in the
 // local database. This is not exhaustive, as the actual list formatting
 // is covered by SafeBrowsingProtocolManagerTest.TestChunkStrings.
 TEST_F(SafeBrowsingProtocolManagerTest, ExistingDatabase) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   std::vector<SBListChunkRanges> ranges;
   SBListChunkRanges range_phish(safe_browsing_util::kPhishingList);
   range_phish.adds = "adds_phish";
   range_phish.subs = "subs_phish";
   ranges.push_back(range_phish);
 
   SBListChunkRanges range_unknown("unknown_list");
   range_unknown.adds = "adds_unknown";
   range_unknown.subs = "subs_unknown";
   ranges.push_back(range_unknown);
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     ranges,
                                     false)));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ASSERT_TRUE(url_fetcher);
   EXPECT_EQ(net::LOAD_DISABLE_CACHE, url_fetcher->GetLoadFlags());
   EXPECT_EQ("goog-phish-shavar;a:adds_phish:s:subs_phish\n"
             "unknown_list;a:adds_unknown:s:subs_unknown\n"
             "goog-malware-shavar;\n",
             url_fetcher->upload_data());
   EXPECT_EQ(GURL("https://prefix.com/foo/downloads?client=unittest&appver=1.0"
                  "&pver=2.2" + key_param_),
             url_fetcher->GetOriginalURL());
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString("");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
-// Tests what happens when there is a response with no chunks.
-TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseEmptyBody) {
+

[mattm, 2012/12/05 23:58:10]

extra line

[cbentzel, 2012/12/17 11:24:10]

Done.

+TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseBadyBody) {

[mattm, 2012/12/05 23:58:10]

s/Bady/Bad/

[cbentzel, 2012/12/17 11:24:10]

Done.

   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
-  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+  EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
 
-  // The update response is successful, but an empty body.
+  // The update response is successful, but an invalid body.
   url_fetcher->set_status(net::URLRequestStatus());
   url_fetcher->set_response_code(200);
-  url_fetcher->SetResponseString("");
+  url_fetcher->SetResponseString("THIS_IS_A_BAD_RESPONSE");
   url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
-TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseBadBody) {
+// Tests what happens when there is an error in the update response.
+TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseHttpError) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
   EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
 
-  // The update response is successful, but an invalid body.
+  // Go ahead and respond to it.
   url_fetcher->set_status(net::URLRequestStatus());
-  url_fetcher->set_response_code(200);
-  url_fetcher->SetResponseString("THIS_IS_A_BAD_RESPONSE");
+  url_fetcher->set_response_code(404);
+  url_fetcher->SetResponseString("");
   url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
-// Tests what happens when there is an error in the update response.
-TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseHttpError) {
+// Tests what happens when there is an error with the connection.
+TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseConnectionError) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
   EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
 
   // Go ahead and respond to it.
-  url_fetcher->set_status(net::URLRequestStatus());
-  url_fetcher->set_response_code(404);
-  url_fetcher->SetResponseString("");
+  url_fetcher->set_status(net::URLRequestStatus(net::URLRequestStatus::FAILED,
+                                                net::ERR_CONNECTION_RESET));
   url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
-// Tests what happens when there is an error with the connection.
-TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseConnectionError) {
+// Tests what happens when there is a timeout before an update response.
+TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseTimeout) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
   EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
 
-  // Go ahead and respond to it.
-  url_fetcher->set_status(net::URLRequestStatus(net::URLRequestStatus::FAILED,
-                                                net::ERR_CONNECTION_RESET));
-  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+  // The first time RunTasks is called above, the update timeout timer is not
+  // handled. This call of RunTasks will handle the update.
+  runner->RunTasks();
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
-// Tests what happens when there is a timeout before an update response.
-TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseTimeout) {
+// Tests what happens when there is a reset command in the response.
+TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseReset) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, ResetDatabase()).Times(1);
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString("r:pleasereset\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
+}
+
+// Tests a single valid update response, followed by a single redirect response
+// that has an valid, but empty body.
+TEST_F(SafeBrowsingProtocolManagerTest, EmptyRedirectResponse) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.

[mattm, 2012/12/06 00:51:40]

s/reset/redirect/
(in following tests too)

[cbentzel, 2012/12/17 11:24:10]

Done.

+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/path\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // We should have another request.
+  net::TestURLFetcher* chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(chunk_url_fetcher != NULL);

[mattm, 2012/12/06 00:51:40]

Should test the fetched URL value as well, at least in one of these tests.

[cbentzel, 2012/12/17 11:24:10]

Done.

+  chunk_url_fetcher->set_status(net::URLRequestStatus());
+  chunk_url_fetcher->set_response_code(200);
+  chunk_url_fetcher->SetResponseString("");
+  chunk_url_fetcher->delegate()->OnURLFetchComplete(chunk_url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
+}
+
+// Tests a single valid update response, followed by a single redirect response
+// that has an invalid body.
+TEST_F(SafeBrowsingProtocolManagerTest, InvalidRedirectResponse) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
   EXPECT_CALL(test_delegate, UpdateFinished(false)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
 
-  // The first time RunTasks is called above, the update timeout timer is not
-  // handled. This call of RunTasks will handle the update.
-  runner->RunTasks();
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/path\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // We should have another request.
+  net::TestURLFetcher* chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(chunk_url_fetcher != NULL);
+  chunk_url_fetcher->set_status(net::URLRequestStatus());
+  chunk_url_fetcher->set_response_code(200);
+  chunk_url_fetcher->SetResponseString("THIS IS AN INVALID RESPONSE");
+  chunk_url_fetcher->delegate()->OnURLFetchComplete(chunk_url_fetcher);
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }
 
-// Tests what happens when there is a reset command in the response.
-TEST_F(SafeBrowsingProtocolManagerTest, UpdateResponseReset) {
+// Tests a single valid update response, followed by a single redirect response
+// containing chunks.
+TEST_F(SafeBrowsingProtocolManagerTest, SingleRedirectResponseWithChunks) {
   scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
       new ImmediateSingleThreadTaskRunner());
   base::ThreadTaskRunnerHandle runner_handler(runner);
+  net::TestURLFetcherFactory url_fetcher_factory;
+
+  testing::StrictMock<MockProtocolDelegate> test_delegate;
+  EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
+  EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
+      Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
+                                    std::vector<SBListChunkRanges>(),
+                                    false)));
+  EXPECT_CALL(test_delegate, AddChunks("goog-phish-shavar", _, _)).WillOnce(
+      Invoke(HandleAddChunks));
+  EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
+
+  scoped_ptr<SafeBrowsingProtocolManager> pm(
+      CreateProtocolManager(&test_delegate));
+
+  // Kick off initialization. This returns chunks from the DB synchronously.
+  pm->ForceScheduleNextUpdate(base::TimeDelta());
+  runner->RunTasks();
+
+  // We should have an URLFetcher at this point in time.
+  net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
+  ValidateUpdateFetcherRequest(url_fetcher);
+
+  // The update response is successful, and has a reset command.
+  url_fetcher->set_status(net::URLRequestStatus());
+  url_fetcher->set_response_code(200);
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/path\n");
+  url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // We should have another request.
+  net::TestURLFetcher* chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(chunk_url_fetcher != NULL);
+  chunk_url_fetcher->set_status(net::URLRequestStatus());
+  chunk_url_fetcher->set_response_code(200);
+  chunk_url_fetcher->SetResponseString(
+      "a:4:4:9\n"
+      "host\1fdaf");
+  chunk_url_fetcher->delegate()->OnURLFetchComplete(chunk_url_fetcher);
+
+  // The AddChunksCallback needs to be invoked
+  runner->RunTasks();
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
+}
+
+// Tests a single valid update response, followed by multiple redirect responses
+// containing chunks.
+TEST_F(SafeBrowsingProtocolManagerTest, MultipleRedirectResponsesWithChunks) {
+  scoped_refptr<ImmediateSingleThreadTaskRunner> runner(
+      new ImmediateSingleThreadTaskRunner());
+  base::ThreadTaskRunnerHandle runner_handler(runner);
   net::TestURLFetcherFactory url_fetcher_factory;
 
   testing::StrictMock<MockProtocolDelegate> test_delegate;
   EXPECT_CALL(test_delegate, UpdateStarted()).Times(1);
   EXPECT_CALL(test_delegate, GetChunks(_)).WillOnce(
       Invoke(testing::CreateFunctor(InvokeGetChunksCallback,
                                     std::vector<SBListChunkRanges>(),
                                     false)));
-  EXPECT_CALL(test_delegate, ResetDatabase()).Times(1);
+  EXPECT_CALL(test_delegate, AddChunks("goog-phish-shavar", _, _)).
+      WillRepeatedly(Invoke(HandleAddChunks));
   EXPECT_CALL(test_delegate, UpdateFinished(true)).Times(1);
 
   scoped_ptr<SafeBrowsingProtocolManager> pm(
       CreateProtocolManager(&test_delegate));
 
   // Kick off initialization. This returns chunks from the DB synchronously.
   pm->ForceScheduleNextUpdate(base::TimeDelta());
   runner->RunTasks();
 
   // We should have an URLFetcher at this point in time.
   net::TestURLFetcher* url_fetcher = url_fetcher_factory.GetFetcherByID(0);
   ValidateUpdateFetcherRequest(url_fetcher);
 
   // The update response is successful, and has a reset command.
   url_fetcher->set_status(net::URLRequestStatus());
   url_fetcher->set_response_code(200);
-  url_fetcher->SetResponseString("r:pleasereset\n");
+  url_fetcher->SetResponseString(
+      "i:goog-phish-shavar\n"
+      "u:redirect-server.example.com/one\n"
+      "u:redirect-server.example.com/two\n");
   url_fetcher->delegate()->OnURLFetchComplete(url_fetcher);
+
+  // First redirect request.
+  net::TestURLFetcher* first_chunk_url_fetcher =

[mattm, 2012/12/06 00:51:40]

should probably check the URLs on these requests too..

[cbentzel, 2012/12/17 11:24:10]

Done.

+      url_fetcher_factory.GetFetcherByID(1);
+  ASSERT_TRUE(first_chunk_url_fetcher != NULL);
+  first_chunk_url_fetcher->set_status(net::URLRequestStatus());
+  first_chunk_url_fetcher->set_response_code(200);
+  first_chunk_url_fetcher->SetResponseString(
+      "a:4:4:9\n"
+      "host\1aaaa");
+  first_chunk_url_fetcher->delegate()->OnURLFetchComplete(
+      first_chunk_url_fetcher);
+
+  // Invoke the AddChunksCallback
+  runner->RunTasks();
+
+  // Second redirect request.
+  net::TestURLFetcher* second_chunk_url_fetcher =
+      url_fetcher_factory.GetFetcherByID(2);
+  ASSERT_TRUE(second_chunk_url_fetcher != NULL);
+  second_chunk_url_fetcher->set_status(net::URLRequestStatus());
+  second_chunk_url_fetcher->set_response_code(200);
+  second_chunk_url_fetcher->SetResponseString(
+      "a:5:4:9\n"
+      "host\1bbbb");
+  second_chunk_url_fetcher->delegate()->OnURLFetchComplete(
+      second_chunk_url_fetcher);
+
+  // The AddChunksCallback needs to be invoked
+  runner->RunTasks();
+
+  EXPECT_TRUE(pm->IsUpdateScheduled());
 }