4 files changed, 25 insertions, 35 deletions

diff --git a/fpdfsdk/fpdfformfill_embeddertest.cpp b/fpdfsdk/fpdfformfill_embeddertest.cpp
index 3a6a7448e5..baf4853f17 100644
--- a/fpdfsdk/fpdfformfill_embeddertest.cpp
+++ b/fpdfsdk/fpdfformfill_embeddertest.cpp
@@ -45,7 +45,7 @@ TEST_F(FPDFFormFillEmbeddertest, BUG_507316) {
   SetDelegate(&delegate);
 
   EXPECT_TRUE(OpenDocument("bug_507316.pdf"));
-  FPDF_PAGE page = LoadAndCachePage(2);
+  FPDF_PAGE page = LoadPage(2);
   EXPECT_TRUE(page);
   DoOpenActions();
   delegate.AdvanceTime(4000);
diff --git a/fpdfsdk/fsdk_baseform_embeddertest.cpp b/fpdfsdk/fsdk_baseform_embeddertest.cpp
index 622f1a8c04..7af8262c55 100644
--- a/fpdfsdk/fsdk_baseform_embeddertest.cpp
+++ b/fpdfsdk/fsdk_baseform_embeddertest.cpp
@@ -26,9 +26,12 @@ class FSDKBaseFormEmbeddertest : public EmbedderTest {};
 
 TEST_F(FSDKBaseFormEmbeddertest, CBA_AnnotIterator) {
   EXPECT_TRUE(OpenDocument("annotiter.pdf"));
-  EXPECT_TRUE(LoadPage(0));
-  EXPECT_TRUE(LoadPage(1));
-  EXPECT_TRUE(LoadPage(2));
+  FPDF_PAGE page0 = LoadPage(0);
+  FPDF_PAGE page1 = LoadPage(1);
+  FPDF_PAGE page2 = LoadPage(2);
+  EXPECT_TRUE(page0);
+  EXPECT_TRUE(page1);
+  EXPECT_TRUE(page2);
 
   CFX_FloatRect LeftBottom(200, 200, 220, 220);
   CFX_FloatRect RightBottom(400, 201, 420, 221);
@@ -112,4 +115,7 @@ TEST_F(FSDKBaseFormEmbeddertest, CBA_AnnotIterator) {
     pAnnot = iter.GetPrevAnnot(pAnnot);
     EXPECT_EQ(iter.GetLastAnnot(), pAnnot);
   }
+  UnloadPage(page2);
+  UnloadPage(page1);
+  UnloadPage(page0);
 }
diff --git a/testing/embedder_test.cpp b/testing/embedder_test.cpp
index b1207fba82..e10e2a81e8 100644
--- a/testing/embedder_test.cpp
+++ b/testing/embedder_test.cpp
@@ -242,21 +242,19 @@ int EmbedderTest::GetPageCount() {
 }
 
 FPDF_PAGE EmbedderTest::LoadPage(int page_number) {
+  // First check whether it is loaded already.
+  auto it = page_map_.find(page_number);
+  if (it != page_map_.end())
+    return it->second;
+
   FPDF_PAGE page = FPDF_LoadPage(document_, page_number);
   if (!page) {
     return nullptr;
   }
   FORM_OnAfterLoadPage(page, form_handle_);
   FORM_DoPageAAction(page, form_handle_, FPDFPAGE_AACTION_OPEN);
-  return page;
-}
-
-FPDF_PAGE EmbedderTest::LoadAndCachePage(int page_number) {
-  FPDF_PAGE page = delegate_->GetPage(form_handle_, document_, page_number);
-  if (!page) {
-    return nullptr;
-  }
-  FORM_DoPageAAction(page, form_handle_, FPDFPAGE_AACTION_OPEN);
+  // Cache the page.
+  page_map_[page_number] = page;
   return page;
 }
 
@@ -278,20 +276,12 @@ void EmbedderTest::UnloadPage(FPDF_PAGE page) {
   FPDF_ClosePage(page);
 }
 
-FPDF_PAGE EmbedderTest::Delegate::GetPage(FPDF_FORMHANDLE form_handle,
+FPDF_PAGE EmbedderTest::Delegate::GetPage(FPDF_FORMFILLINFO* info,
                                           FPDF_DOCUMENT document,
                                           int page_index) {
-  auto it = m_pageMap.find(page_index);
-  if (it != m_pageMap.end()) {
-    return it->second;
-  }
-  FPDF_PAGE page = FPDF_LoadPage(document, page_index);
-  if (!page) {
-    return nullptr;
-  }
-  m_pageMap[page_index] = page;
-  FORM_OnAfterLoadPage(page, form_handle);
-  return page;
+  EmbedderTest* test = static_cast<EmbedderTest*>(info);
+  auto it = test->page_map_.find(page_index);
+  return it != test->page_map_.end() ? it->second : nullptr;
 }
 
 // static
@@ -329,8 +319,8 @@ void EmbedderTest::KillTimerTrampoline(FPDF_FORMFILLINFO* info, int id) {
 FPDF_PAGE EmbedderTest::GetPageTrampoline(FPDF_FORMFILLINFO* info,
                                           FPDF_DOCUMENT document,
                                           int page_index) {
-  EmbedderTest* test = static_cast<EmbedderTest*>(info);
-  return test->delegate_->GetPage(test->form_handle(), document, page_index);
+  return static_cast<EmbedderTest*>(info)->delegate_->GetPage(info, document,
+                                                              page_index);
 }
 
 // Can't use gtest-provided main since we need to stash the path to the
diff --git a/testing/embedder_test.h b/testing/embedder_test.h
index 6b814a7267..153ca6e5f3 100644
--- a/testing/embedder_test.h
+++ b/testing/embedder_test.h
@@ -51,12 +51,9 @@ class EmbedderTest : public ::testing::Test,
     virtual void KillTimer(int id) {}
 
     // Equivalent to FPDF_FORMFILLINFO::FFI_GetPage().
-    virtual FPDF_PAGE GetPage(FPDF_FORMHANDLE form_handle,
+    virtual FPDF_PAGE GetPage(FPDF_FORMFILLINFO* info,
                               FPDF_DOCUMENT document,
                               int page_index);
-
-   private:
-    std::map<int, FPDF_PAGE> m_pageMap;
   };
 
   EmbedderTest();
@@ -101,10 +98,6 @@ class EmbedderTest : public ::testing::Test,
   // Load a specific page of the open document.
   virtual FPDF_PAGE LoadPage(int page_number);
 
-  // Load a specific page of the open document using delegate_->GetPage.
-  // delegate_->GetPage also caches loaded page.
-  virtual FPDF_PAGE LoadAndCachePage(int page_number);
-
   // Convert a loaded page into a bitmap.
   virtual FPDF_BITMAP RenderPage(FPDF_PAGE page);
 
@@ -130,6 +123,7 @@ class EmbedderTest : public ::testing::Test,
   TestLoader* loader_;
   size_t file_length_;
   std::unique_ptr<char, pdfium::FreeDeleter> file_contents_;
+  std::map<int, FPDF_PAGE> page_map_;
 
  private:
   static void UnsupportedHandlerTrampoline(UNSUPPORT_INFO*, int type);