diff --git a/contrib/queries/src/test/org/apache/lucene/search/trie/TestIntTrieTokenStream.java b/contrib/queries/src/test/org/apache/lucene/search/trie/TestIntTrieTokenStream.java
index cdd4f17daad..e53b6e22d1e 100644
--- a/contrib/queries/src/test/org/apache/lucene/search/trie/TestIntTrieTokenStream.java
+++ b/contrib/queries/src/test/org/apache/lucene/search/trie/TestIntTrieTokenStream.java
@@ -29,8 +29,10 @@ public class TestIntTrieTokenStream extends LuceneTestCase {
   public void testStreamNewAPI() throws Exception {
     final IntTrieTokenStream stream=new IntTrieTokenStream(value, precisionStep);
     stream.setUseNewAPI(true);
-    final ShiftAttribute shiftAtt = (ShiftAttribute) stream.addAttribute(ShiftAttribute.class);
-    final TermAttribute termAtt = (TermAttribute) stream.addAttribute(TermAttribute.class);
+    final ShiftAttribute shiftAtt = (ShiftAttribute) stream.getAttribute(ShiftAttribute.class);
+    assertNotNull("Has shift attribute", shiftAtt);
+    final TermAttribute termAtt = (TermAttribute) stream.getAttribute(TermAttribute.class);
+    assertNotNull("Has term attribute", termAtt);
     for (int shift=0; shift<32; shift+=precisionStep) {
       assertTrue("New token is available", stream.incrementToken());
       assertEquals("Shift value", shift, shiftAtt.getShift());
diff --git a/contrib/queries/src/test/org/apache/lucene/search/trie/TestLongTrieTokenStream.java b/contrib/queries/src/test/org/apache/lucene/search/trie/TestLongTrieTokenStream.java
index f2341220f99..e116f9eca66 100644
--- a/contrib/queries/src/test/org/apache/lucene/search/trie/TestLongTrieTokenStream.java
+++ b/contrib/queries/src/test/org/apache/lucene/search/trie/TestLongTrieTokenStream.java
@@ -29,8 +29,10 @@ public class TestLongTrieTokenStream extends LuceneTestCase {
   public void testStreamNewAPI() throws Exception {
     final LongTrieTokenStream stream=new LongTrieTokenStream(value, precisionStep);
     stream.setUseNewAPI(true);
-    final ShiftAttribute shiftAtt = (ShiftAttribute) stream.addAttribute(ShiftAttribute.class);
-    final TermAttribute termAtt = (TermAttribute) stream.addAttribute(TermAttribute.class);
+    final ShiftAttribute shiftAtt = (ShiftAttribute) stream.getAttribute(ShiftAttribute.class);
+    assertNotNull("Has shift attribute", shiftAtt);
+    final TermAttribute termAtt = (TermAttribute) stream.getAttribute(TermAttribute.class);
+    assertNotNull("Has term attribute", termAtt);
     for (int shift=0; shift<64; shift+=precisionStep) {
       assertTrue("New token is available", stream.incrementToken());
       assertEquals("Shift value", shift, shiftAtt.getShift());