Store page no for distilled pages undergoing distillation.

components/dom_distiller/core/distiller_unittest.cc

Index: components/dom_distiller/core/distiller_unittest.cc
diff --git a/components/dom_distiller/core/distiller_unittest.cc b/components/dom_distiller/core/distiller_unittest.cc
index 3865aefa4d60fbcec500aa27e4df8a1fdf89fb8f..296242fceed03d1702633185a0da086ca4828066 100644
--- a/components/dom_distiller/core/distiller_unittest.cc
+++ b/components/dom_distiller/core/distiller_unittest.cc
@@ -28,6 +28,8 @@ using::testing::Return;
 using::testing::_;
 
 namespace {
+const size_t kMaxPagesInArticle = 10;
+
   const char kTitle[] = "Title";
   const char kContent[] = "Content";
   const char kURL[] = "http://a.com/";
@@ -56,6 +58,21 @@ namespace {
     return list.Pass();
   }
 
+  class TestDistillerWithMaxPageLimit : public dom_distiller::DistillerImpl {
+   public:
+    TestDistillerWithMaxPageLimit(
+        const dom_distiller::DistillerPageFactory& distiller_page_factory,
+        const dom_distiller::DistillerURLFetcherFactory&
+            distiller_url_fetcher_factory)
+        : DistillerImpl(distiller_page_factory, distiller_url_fetcher_factory) {
+    }
+
+   protected:
+    virtual size_t GetMaxNumPagesInArticle() const OVERRIDE {
+      return kMaxPagesInArticle;
+    }
+  };
+
 }  // namespace
 
 namespace dom_distiller {
@@ -117,7 +134,6 @@ class MockDistillerPageFactory : public DistillerPageFactory {
   }
 };
 
-
 class DistillerTest : public testing::Test {
  public:
   virtual ~DistillerTest() {}
@@ -290,4 +306,94 @@ TEST_F(DistillerTest, DistillLinkLoop) {
   EXPECT_EQ(article_proto_->pages_size(), 1);
 }
 
+TEST_F(DistillerTest, CheckMaxPageLimit) {
+  base::MessageLoopForUI loop;
+  string page_urls[kMaxPagesInArticle];
+  scoped_ptr<base::ListValue> list[kMaxPagesInArticle];
+
+  // Note: Next page url of the last page of article is set. So distiller will
+  // try to do kMaxPagesInArticle + 1 calls if the max article limit does not
+  // work.
+  string url_prefix = "http://a.com/";
+  for (size_t page_num = 0; page_num < kMaxPagesInArticle; ++page_num) {
+    page_urls[page_num] = url_prefix + base::IntToString(page_num + 1);
+    string content = "Content for page:" + base::IntToString(page_num);
+    string next_page_url = url_prefix + base::IntToString(page_num + 2);
+    list[page_num] = CreateDistilledValueReturnedFromJS(
+        kTitle, content, vector<int>(), next_page_url);
+  }
+
+  EXPECT_CALL(page_factory_, CreateDistillerPageMock(_))
+      .WillOnce(CreateMockDistillerPages(
+          list, page_urls, static_cast<int>(kMaxPagesInArticle)));
+
+  distiller_.reset(
+      new TestDistillerWithMaxPageLimit(page_factory_, url_fetcher_factory_));
+  distiller_->Init();
+  distiller_->DistillPage(
+      GURL(page_urls[0]),
+      base::Bind(&DistillerTest::OnDistillPageDone, base::Unretained(this)));
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ(kTitle, article_proto_->title());
+  EXPECT_EQ(kMaxPagesInArticle,
+            static_cast<size_t>(article_proto_->pages_size()));
+}
+
+TEST_F(DistillerTest, SinglePageDistillationFailure) {
+  base::MessageLoopForUI loop;
+  // To simulate failure set the value returned as
+  scoped_ptr<base::Value> emptyValue(base::Value::CreateNullValue());
+  EXPECT_CALL(page_factory_, CreateDistillerPageMock(_))
+      .WillOnce(CreateMockDistillerPage(emptyValue.get(), GURL(kURL)));
+  distiller_.reset(new DistillerImpl(page_factory_, url_fetcher_factory_));
+  distiller_->Init();
+  distiller_->DistillPage(
+      GURL(kURL),
+      base::Bind(&DistillerTest::OnDistillPageDone, base::Unretained(this)));
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ("", article_proto_->title());
+  EXPECT_EQ(0, article_proto_->pages_size());
+}
+
+TEST_F(DistillerTest, MultiplePagesDistillationFailure) {
+  base::MessageLoopForUI loop;
+  const int kNumPages = 8;
+  string content[kNumPages];
+  string page_urls[kNumPages];
+  scoped_ptr<base::Value> distilled_values[kNumPages];
+  // The page number of the failed page.
+  int failed_page_no = 3;
+  string url_prefix = "http://a.com/";
+  for (int page_no = 0; page_no < kNumPages; ++page_no) {
+    page_urls[page_no] = url_prefix + base::IntToString(page_no);
+    content[page_no] = "Content for page:" + base::IntToString(page_no);
+    string next_page_url = url_prefix + base::IntToString(page_no + 1);
+    if (page_no != failed_page_no) {
+      distilled_values[page_no] = CreateDistilledValueReturnedFromJS(
+          kTitle, content[page_no], vector<int>(), next_page_url);
+    } else {
+      distilled_values[page_no].reset(base::Value::CreateNullValue());
+    }
+  }
+
+  // Expect only calls till the failed page number.
+  EXPECT_CALL(page_factory_, CreateDistillerPageMock(_))
+      .WillOnce(CreateMockDistillerPages(
+          distilled_values, page_urls, failed_page_no + 1));
+
+  distiller_.reset(new DistillerImpl(page_factory_, url_fetcher_factory_));
+  distiller_->Init();
+  distiller_->DistillPage(
+      GURL(page_urls[0]),
+      base::Bind(&DistillerTest::OnDistillPageDone, base::Unretained(this)));
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ(kTitle, article_proto_->title());
+  EXPECT_EQ(article_proto_->pages_size(), failed_page_no);
+  for (int page_no = 0; page_no < failed_page_no; ++page_no) {
+    const DistilledPageProto& page = article_proto_->pages(page_no);
+    EXPECT_EQ(content[page_no], page.html());
+    EXPECT_EQ(page_urls[page_no], page.url());
+  }
+}
+
 }  // namespace dom_distiller