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

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

 // Copyright 2016 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.
 
 package org.chromium.chrome.browser.ntp.cards;
 
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ skipping 467 matching lines @@
         // Indices in section are off-by-one (index 0 is the header).
         bindViewHolders(section, 1, 2);
 
         List<SnippetArticle> newSnippets =
                 createDummySuggestions(3, TEST_CATEGORY_ID, "new");
         // Copy the list when passing to the section - it may alter it but we later need it.
         section.updateSuggestions(createSourceFor(new ArrayList<>(newSnippets)));
         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());
     }
 
     /**
      * Tests that the UI does not update the first two items of the section if they have been
      * viewed.
      */
     @Test
     @Feature({"Ntp"})
     public void testUpdateSectionDoesNotReplaceFirstTwoSuggestionWhenSeen() {
         List<SnippetArticle> snippets = createDummySuggestions(4, TEST_CATEGORY_ID, "old");
         // Copy the list when passing to the section - it may alter it but we later need it.
         SuggestionsSection section =
                 createSectionWithSuggestions(new ArrayList<>(snippets));
         assertEquals(4, section.getSuggestionsCount());
 
         // Bind the first two suggestions - indicate that they are being viewed.
         // Indices in section are off-by-one (index 0 is the header).
         bindViewHolders(section, 1, 3);
 
         List<SnippetArticle> newSnippets =
                 createDummySuggestions(3, TEST_CATEGORY_ID, "new");
         // Copy the list when passing to the section - it may alter it but we later need it.
         section.updateSuggestions(createSourceFor(new ArrayList<>(newSnippets)));
         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());
     }
 
     /**
      * Tests that the UI does not update any items of the section if the new list is shorter than
      * what has been viewed.
      */
     @Test
     @Feature({"Ntp"})
     public void testUpdateSectionDoesNothingWhenNewListIsShorterThanItemsSeen() {
         List<SnippetArticle> snippets = createDummySuggestions(4, TEST_CATEGORY_ID, "old");
         // Copy the list when passing to the section - it may alter it but we later need it.
         SuggestionsSection section =
                 createSectionWithSuggestions(new ArrayList<>(snippets));
         assertEquals(4, section.getSuggestionsCount());
 
         // Bind the first two suggestions - indicate that they are being viewed.
         // Indices in section are off-by-one (index 0 is the header).
         bindViewHolders(section, 1, 3);
 
         section.updateSuggestions(createSourceFor(createDummySuggestions(1, TEST_CATEGORY_ID)));
         // Even though the new list has just one suggestion, we need to keep the two seen ones
         // around.
         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());
     }
 
     /**
      * Tests that the UI does not update any items of the section if the current list is shorter
      * than what has been viewed.
      */
     @Test
     @Feature({"Ntp"})
     public void testUpdateSectionDoesNothingWhenCurrentListIsShorterThanItemsSeen() {
         List<SnippetArticle> snippets = createDummySuggestions(3, TEST_CATEGORY_ID, "old");
         // Copy the list when passing to the section - it may alter it but we later need it.
         SuggestionsSection section =
                 createSectionWithSuggestions(new ArrayList<>(snippets));
         assertEquals(3, section.getSuggestionsCount());
 
         // Bind the first two suggestions - indicate that they are being viewed.
         // Indices in section are off-by-one (index 0 is the header).
         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());
 
         section.updateSuggestions(createSourceFor(createDummySuggestions(4, TEST_CATEGORY_ID)));
         // We do not touch the current list if all has been seen.
         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());
     }
 
     /**
      * Tests that the UI does not update when the section has been viewed.
      */
     @Test
     @Feature({"Ntp"})
     public void testUpdateSectionDoesNothingWhenAllSeen() {
         List<SnippetArticle> snippets = createDummySuggestions(4, TEST_CATEGORY_ID, "old");
         SuggestionsSection section = createSectionWithSuggestions(snippets);
         assertEquals(4, section.getSuggestionsCount());
 
         // Bind all the suggestions - indicate that they are being viewed.
         bindViewHolders(section);
 
         section.updateSuggestions(createSourceFor(createDummySuggestions(3, TEST_CATEGORY_ID)));
         verify(mParent, never()).onItemRangeRemoved(any(TreeNode.class), anyInt(), anyInt());
         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());
     }
 
     /**
      * Tests that the UI does not update when anything has been appended.
      */
     @Test
     @Feature({"Ntp"})
     public void testUpdateSectionDoesNothingWhenUserAppended() {
         List<SnippetArticle> snippets = createDummySuggestions(4, TEST_CATEGORY_ID, "old");
         SuggestionsSection section = createSectionWithSuggestions(snippets);
 
         // Append another 3 suggestions.
         List<SnippetArticle> appendedSnippets =
                 createDummySuggestions(3, TEST_CATEGORY_ID, "appended");
         section.appendSuggestions(appendedSnippets, /* userRequested = */ true);
 
         // 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 =
                 createDummySuggestions(5, TEST_CATEGORY_ID, "new");
         section.updateSuggestions(createSourceFor(newSnippets));
 
         // All previous snippets should be in place.
-        verifySnippets(section, snippets);
+        assertEquals(snippets, getSuggestions(section));
 
         assertTrue(section.isDataStale());
     }
 
     @Test
     @Feature({"Ntp"})
     public void testCardIsNotifiedWhenBecomingFirst() {
         List<SnippetArticle> suggestions = createDummySuggestions(5, /* categoryId = */ 42);
         SuggestionsSection section = createSectionWithFetchAction(false);
         section.appendSuggestions(suggestions, /* userRequested = */ false);
@@ skipping 86 matching lines @@
         return section;
     }
 
     @SafeVarargs
     private static <T> Set<T> setOf(T... elements) {
         Set<T> set = new TreeSet<>();
         set.addAll(Arrays.asList(elements));
         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);
         return createSection(builder.build());
     }
 
     private SuggestionsSection createSection(SuggestionsCategoryInfo info) {
         SuggestionsSection section = new SuggestionsSection(
                 mDelegate, mUiDelegate, mock(SuggestionsRanker.class), mBridge, info);
         section.setParent(mParent);
@@ skipping 27 matching lines @@
             section.getActionItemForTesting().performAction(manager);
         }
 
         verify(section.getCategoryInfo(),
                 (action == ContentSuggestionsAdditionalAction.VIEW_ALL ? times(1) : never()))
                 .performViewAllAction(navDelegate);
         verify(suggestionsSource,
                 (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++));
-        }
-    }
 }