[Suggestions] Remove TreeNode.getSuggestionAt() in favor of a visitor.

chrome/android/junit/src/org/chromium/chrome/browser/ntp/cards/SuggestionsSectionTest.java

Index: chrome/android/junit/src/org/chromium/chrome/browser/ntp/cards/SuggestionsSectionTest.java
diff --git a/chrome/android/junit/src/org/chromium/chrome/browser/ntp/cards/SuggestionsSectionTest.java b/chrome/android/junit/src/org/chromium/chrome/browser/ntp/cards/SuggestionsSectionTest.java
index 1d421c8243f0389836c4095f6f59e7c21febee61..3eecc79d87fa18344592c9afe4dcfc9e74ca572a 100644
--- a/chrome/android/junit/src/org/chromium/chrome/browser/ntp/cards/SuggestionsSectionTest.java
+++ b/chrome/android/junit/src/org/chromium/chrome/browser/ntp/cards/SuggestionsSectionTest.java
@@ -485,11 +485,12 @@ public class SuggestionsSectionTest {
         verify(mParent).onItemRangeRemoved(section, 2, 3);
         verify(mParent).onItemRangeInserted(section, 2, 2);
         assertEquals(3, section.getSuggestionsCount());
-        assertEquals(snippets.get(0), section.getSuggestionAt(1));
-        assertNotEquals(snippets.get(1), section.getSuggestionAt(2));
-        assertEquals(newSnippets.get(0), section.getSuggestionAt(2));
-        assertNotEquals(snippets.get(2), section.getSuggestionAt(3));
-        assertEquals(newSnippets.get(1), section.getSuggestionAt(3));
+        List<SnippetArticle> sectionSuggestions = getSuggestions(section);
+        assertEquals(snippets.get(0), sectionSuggestions.get(0));
+        assertNotEquals(snippets.get(1), sectionSuggestions.get(1));
+        assertEquals(newSnippets.get(0), sectionSuggestions.get(1));
+        assertNotEquals(snippets.get(2), sectionSuggestions.get(2));
+        assertEquals(newSnippets.get(1), sectionSuggestions.get(2));
 
         assertTrue(section.isDataStale());
     }
@@ -518,10 +519,11 @@ public class SuggestionsSectionTest {
         verify(mParent).onItemRangeRemoved(section, 3, 2);
         verify(mParent).onItemRangeInserted(section, 3, 1);
         assertEquals(3, section.getSuggestionsCount());
-        assertEquals(snippets.get(0), section.getSuggestionAt(1));
-        assertEquals(snippets.get(1), section.getSuggestionAt(2));
-        assertNotEquals(snippets.get(2), section.getSuggestionAt(3));
-        assertEquals(newSnippets.get(0), section.getSuggestionAt(3));
+        List<SnippetArticle> sectionSuggestions = getSuggestions(section);
+        assertEquals(snippets.get(0), sectionSuggestions.get(0));
+        assertEquals(snippets.get(1), sectionSuggestions.get(1));
+        assertNotEquals(snippets.get(2), sectionSuggestions.get(2));
+        assertEquals(newSnippets.get(0), sectionSuggestions.get(2));
 
         assertTrue(section.isDataStale());
     }
@@ -549,8 +551,9 @@ public class SuggestionsSectionTest {
         verify(mParent).onItemRangeRemoved(section, 3, 2);
         verify(mParent, never()).onItemRangeInserted(any(TreeNode.class), anyInt(), anyInt());
         assertEquals(2, section.getSuggestionsCount());
-        assertEquals(snippets.get(0), section.getSuggestionAt(1));
-        assertEquals(snippets.get(1), section.getSuggestionAt(2));
+        List<SnippetArticle> sectionSuggestions = getSuggestions(section);
+        assertEquals(snippets.get(0), sectionSuggestions.get(0));
+        assertEquals(snippets.get(1), sectionSuggestions.get(1));
 
         assertTrue(section.isDataStale());
     }
@@ -573,8 +576,9 @@ public class SuggestionsSectionTest {
         bindViewHolders(section, 1, 3);
 
         // Remove last two items.
-        section.removeSuggestionById(section.getSuggestionAt(3).mIdWithinCategory);
-        section.removeSuggestionById(section.getSuggestionAt(2).mIdWithinCategory);
+        List<SnippetArticle> sectionSuggestions = getSuggestions(section);
+        section.removeSuggestionById(sectionSuggestions.get(2).mIdWithinCategory);
+        section.removeSuggestionById(sectionSuggestions.get(1).mIdWithinCategory);
         reset(mParent);
 
         assertEquals(1, section.getSuggestionsCount());
@@ -584,7 +588,7 @@ public class SuggestionsSectionTest {
         verify(mParent, never()).onItemRangeRemoved(any(TreeNode.class), anyInt(), anyInt());
         verify(mParent, never()).onItemRangeInserted(any(TreeNode.class), anyInt(), anyInt());
         assertEquals(1, section.getSuggestionsCount());
-        assertEquals(snippets.get(0), section.getSuggestionAt(1));
+        assertEquals(snippets.get(0), sectionSuggestions.get(0));
 
         assertTrue(section.isDataStale());
     }
@@ -607,7 +611,7 @@ public class SuggestionsSectionTest {
         verify(mParent, never()).onItemRangeInserted(any(TreeNode.class), anyInt(), anyInt());
 
         // All old snippets should be in place.
-        verifySnippets(section, snippets);
+        assertEquals(snippets, getSuggestions(section));
 
         assertTrue(section.isDataStale());
     }
@@ -628,7 +632,7 @@ public class SuggestionsSectionTest {
 
         // All 7 snippets should be in place.
         snippets.addAll(appendedSnippets);
-        verifySnippets(section, snippets);
+        assertEquals(snippets, getSuggestions(section));
 
         // Try to replace them with another list. Should have no effect.
         List<SnippetArticle> newSnippets =
@@ -636,7 +640,7 @@ public class SuggestionsSectionTest {
         section.updateSuggestions(createSourceFor(newSnippets));
 
         // All previous snippets should be in place.
-        verifySnippets(section, snippets);
+        assertEquals(snippets, getSuggestions(section));
 
         assertTrue(section.isDataStale());
     }
@@ -743,6 +747,17 @@ public class SuggestionsSectionTest {
         return set;
     }
 
+    private static List<SnippetArticle> getSuggestions(TreeNode root) {
+        final List<SnippetArticle> suggestions = new ArrayList<>();
+        root.visitItems(new NodeVisitor() {
+            @Override
+            public void visitSuggestion(SnippetArticle suggestion) {
+                suggestions.add(suggestion);
+            }
+        });
+        return suggestions;
+    }
+
     private SuggestionsSection createSectionWithFetchAction(boolean hasFetchAction) {
         CategoryInfoBuilder builder = new CategoryInfoBuilder(TEST_CATEGORY_ID).showIfEmpty();
         if (hasFetchAction) builder.withAction(ContentSuggestionsAdditionalAction.FETCH);
@@ -790,13 +805,4 @@ public class SuggestionsSectionTest {
                 (action == ContentSuggestionsAdditionalAction.FETCH ? times(1) : never()))
                 .fetchSuggestions(anyInt(), any(String[].class));
     }
-
-    private static void verifySnippets(SuggestionsSection section, List<SnippetArticle> snippets) {
-        assertEquals(snippets.size(), section.getSuggestionsCount());
-        // Indices in section are off-by-one (index 0 is the header).
-        int index = 1;
-        for (SnippetArticle snippet : snippets) {
-            assertEquals(snippet, section.getSuggestionAt(index++));
-        }
-    }
 }