JNI refactor: @NativeMethods conversion for Contextual Search.

This CL was partially created by
//base/android/jni_generator/jni_refactorer.py.

The tests were refactored manually to use JniMocker.

This allows us to get rid of ContextualSearchContextForTest.

Bug: 929661
Change-Id: I5dcfc91df003d26cefad9aaa2bb2bf0b2afcabf9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1817117
Commit-Queue: Eric Stevenson <estevenson@chromium.org>
Reviewed-by: Donn Denman <donnd@chromium.org>
Cr-Commit-Position: refs/heads/master@{#699509}

chrome/android/chrome_junit_test_java_sources.gni

@@ -41,7 +41,6 @@ chrome_junit_test_java_sources = [
   "junit/src/org/chromium/chrome/browser/compositor/overlays/strip/StripLayoutHelperTest.java",
   "junit/src/org/chromium/chrome/browser/contextmenu/ChromeContextMenuPopulatorTest.java",
   "junit/src/org/chromium/chrome/browser/contextmenu/RevampedContextMenuCoordinatorTest.java",
-  "junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchContextForTest.java",
   "junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchContextTest.java",
   "junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchEntityHeuristicTest.java",
   "junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchInternalStateTest.java",

chrome/android/java/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchContext.java

@@ -11,6 +11,7 @@ import androidx.annotation.Nullable;
 import org.chromium.base.Log;
 import org.chromium.base.VisibleForTesting;
 import org.chromium.base.annotations.CalledByNative;
+import org.chromium.base.annotations.NativeMethods;
 import org.chromium.content_public.browser.WebContents;
 
 /**
@@ -78,7 +79,7 @@ public abstract class ContextualSearchContext {
     private int mPreviousUserInteractions;
 
     /** A {@link ContextualSearchContext} that ignores changes to the selection. */
-    private static class ChangeIgnoringContext extends ContextualSearchContext {
+    static class ChangeIgnoringContext extends ContextualSearchContext {
         @Override
         void onSelectionChanged() {}
     }
@@ -91,7 +92,7 @@ public abstract class ContextualSearchContext {
      * @return A {@link ContextualSearchContext} or {@code null} if the insertion point happens to
      *         miss a word (e.g. it has non-word characters on both sides).
      */
-    static public @Nullable ContextualSearchContext getContextForInsertionPoint(
+    public static @Nullable ContextualSearchContext getContextForInsertionPoint(
             String surroundingText, int insertionPointOffset) {
         ContextualSearchContext context = new ChangeIgnoringContext();
         context.setSurroundingText(
@@ -117,7 +118,7 @@ public abstract class ContextualSearchContext {
      * Constructs a context that tracks the selection and some amount of page content.
      */
     ContextualSearchContext() {
-        mNativePointer = nativeInit();
+        mNativePointer = ContextualSearchContextJni.get().init(this);
         mHasSetResolveProperties = false;
     }
 
@@ -137,17 +138,18 @@ public abstract class ContextualSearchContext {
         mHomeCountry = homeCountry;
         mPreviousEventId = previousEventId;
         mPreviousUserInteractions = previousUserInteractions;
-        nativeSetResolveProperties(getNativePointer(), homeCountry, maySendBasePageUrl,
-                previousEventId, previousUserInteractions);
+        ContextualSearchContextJni.get().setResolveProperties(getNativePointer(), this, homeCountry,
+                maySendBasePageUrl, previousEventId, previousUserInteractions);
     }
 
     /**
      * This method should be called to clean up storage when an instance of this class is
-     * no longer in use.  The nativeDestroy will call the destructor on the native instance.
+     * no longer in use.  The ContextualSearchContextJni.get().destroy will call the destructor on
+     * the native instance.
      */
     public void destroy() {
         assert mNativePointer != 0;
-        nativeDestroy(mNativePointer);
+        ContextualSearchContextJni.get().destroy(mNativePointer, this);
         mNativePointer = 0;
 
         // Also zero out private data that may be sizable.
@@ -191,8 +193,8 @@ public abstract class ContextualSearchContext {
             onSelectionChanged();
         }
         if (setNative) {
-            nativeSetContent(getNativePointer(), mSurroundingText, mSelectionStartOffset,
-                    mSelectionEndOffset);
+            ContextualSearchContextJni.get().setContent(getNativePointer(), this, mSurroundingText,
+                    mSelectionStartOffset, mSelectionEndOffset);
         }
     }
 
@@ -283,7 +285,7 @@ public abstract class ContextualSearchContext {
         mSurroundingText = mInitialSelectedWord;
         mSelectionStartOffset = 0;
         mSelectionEndOffset = mSurroundingText.length();
-        nativeRestrictResolve(mNativePointer);
+        ContextualSearchContextJni.get().restrictResolve(mNativePointer, this);
     }
 
     /**
@@ -298,7 +300,8 @@ public abstract class ContextualSearchContext {
         mSelectionStartOffset += startAdjust;
         mSelectionEndOffset += endAdjust;
         updateInitialSelectedWord();
-        nativeAdjustSelection(getNativePointer(), startAdjust, endAdjust);
+        ContextualSearchContextJni.get().adjustSelection(
+                getNativePointer(), this, startAdjust, endAdjust);
         // Notify of changes.
         onSelectionChanged();
     }
@@ -340,7 +343,8 @@ public abstract class ContextualSearchContext {
     String getDetectedLanguage() {
         assert mSurroundingText != null;
         if (mDetectedLanguage == null) {
-            mDetectedLanguage = nativeDetectLanguage(mNativePointer);
+            mDetectedLanguage =
+                    ContextualSearchContextJni.get().detectLanguage(mNativePointer, this);
         }
         return mDetectedLanguage;
     }
@@ -567,25 +571,18 @@ public abstract class ContextualSearchContext {
         return mNativePointer;
     }
 
-    // ============================================================================================
-    // Native methods.
-    // ============================================================================================
-    @VisibleForTesting
-    protected native long nativeInit();
-    @VisibleForTesting
-    protected native void nativeDestroy(long nativeContextualSearchContext);
-    @VisibleForTesting
-    protected native void nativeSetResolveProperties(long nativeContextualSearchContext,
-            String homeCountry, boolean maySendBasePageUrl, long previousEventId,
-            int previousEventResults);
-    @VisibleForTesting
-    protected native void nativeAdjustSelection(
-            long nativeContextualSearchContext, int startAdjust, int endAdjust);
-    @VisibleForTesting
-    protected native void nativeSetContent(long nativeContextualSearchContext, String content,
-            int selectionStart, int selectionEnd);
-    @VisibleForTesting
-    protected native String nativeDetectLanguage(long nativeContextualSearchContext);
-    @VisibleForTesting
-    protected native void nativeRestrictResolve(long nativeContextualSearchContext);
+    @NativeMethods
+    interface Natives {
+        long init(ContextualSearchContext caller);
+        void destroy(long nativeContextualSearchContext, ContextualSearchContext caller);
+        void setResolveProperties(long nativeContextualSearchContext,
+                ContextualSearchContext caller, String homeCountry, boolean maySendBasePageUrl,
+                long previousEventId, int previousEventResults);
+        void adjustSelection(long nativeContextualSearchContext, ContextualSearchContext caller,
+                int startAdjust, int endAdjust);
+        void setContent(long nativeContextualSearchContext, ContextualSearchContext caller,
+                String content, int selectionStart, int selectionEnd);
+        String detectLanguage(long nativeContextualSearchContext, ContextualSearchContext caller);
+        void restrictResolve(long nativeContextualSearchContext, ContextualSearchContext caller);
+    }
 }

chrome/android/java/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchManager.java

@@ -19,6 +19,7 @@ import org.chromium.base.SysUtils;
 import org.chromium.base.TimeUtils;
 import org.chromium.base.VisibleForTesting;
 import org.chromium.base.annotations.CalledByNative;
+import org.chromium.base.annotations.NativeMethods;
 import org.chromium.base.metrics.RecordUserAction;
 import org.chromium.chrome.R;
 import org.chromium.chrome.browser.ChromeActivity;
@@ -250,7 +251,7 @@ public class ContextualSearchManager
      * @param parentView The parent view to attach Contextual Search UX to.
      */
     public void initialize(ViewGroup parentView) {
-        mNativeContextualSearchManagerPtr = nativeInit();
+        mNativeContextualSearchManagerPtr = ContextualSearchManagerJni.get().init(this);
 
         mParentView = parentView;
         mParentView.getViewTreeObserver().addOnGlobalFocusChangeListener(mOnFocusChangeListener);
@@ -279,7 +280,7 @@ public class ContextualSearchManager
 
         hideContextualSearch(StateChangeReason.UNKNOWN);
         mParentView.getViewTreeObserver().removeOnGlobalFocusChangeListener(mOnFocusChangeListener);
-        nativeDestroy(mNativeContextualSearchManagerPtr);
+        ContextualSearchManagerJni.get().destroy(mNativeContextualSearchManagerPtr, this);
         stopListeningForHideNotifications();
         mTabRedirectHandler.clear();
         mInternalStateController.enter(InternalState.UNDEFINED);
@@ -488,8 +489,8 @@ public class ContextualSearchManager
         WebContents baseWebContents = getBaseWebContents();
         if (baseWebContents != null && mContext != null && mContext.canResolve()) {
             if (isRestrictedResolve) mContext.setRestrictedResolve();
-            nativeStartSearchTermResolutionRequest(
-                    mNativeContextualSearchManagerPtr, mContext, getBaseWebContents());
+            ContextualSearchManagerJni.get().startSearchTermResolutionRequest(
+                    mNativeContextualSearchManagerPtr, this, mContext, getBaseWebContents());
         } else {
             // Something went wrong and we couldn't resolve.
             hideContextualSearch(StateChangeReason.UNKNOWN);
@@ -609,11 +610,10 @@ public class ContextualSearchManager
     }
 
     /**
-     * Called in response to the
-     * {@link ContextualSearchManager#nativeStartSearchTermResolutionRequest} method.
-     * If {@code nativeStartSearchTermResolutionRequest} is called with a previous request sill
-     * pending our native delegate is supposed to cancel all previous requests.  So this code
-     * should only be called with data corresponding to the most recent request.
+     * Called in response to the {@link ContextualSearchManagerJni#startSearchTermResolutionRequest}
+     * method. If {@code startSearchTermResolutionRequest} is called with a previous request sill
+     * pending our native delegate is supposed to cancel all previous requests.  So this code should
+     * only be called with data corresponding to the most recent request.
      * @param isNetworkUnavailable Indicates if the network is unavailable, in which case all other
      *        parameters should be ignored.
      * @param responseCode The HTTP response code. If the code is not OK, the query should be
@@ -787,7 +787,8 @@ public class ContextualSearchManager
         mLoadedSearchUrlTimeMs = System.currentTimeMillis();
         mLastSearchRequestLoaded = mSearchRequest;
         String searchUrl = mSearchRequest.getSearchUrl();
-        nativeWhitelistContextualSearchJsApiUrl(mNativeContextualSearchManagerPtr, searchUrl);
+        ContextualSearchManagerJni.get().whitelistContextualSearchJsApiUrl(
+                mNativeContextualSearchManagerPtr, this, searchUrl);
         mSearchPanel.loadUrlInPanel(searchUrl);
         mDidStartLoadingResolvedSearchRequest = true;
 
@@ -948,12 +949,14 @@ public class ContextualSearchManager
 
     @Override
     public String getAcceptLanguages() {
-        return nativeGetAcceptLanguages(mNativeContextualSearchManagerPtr);
+        return ContextualSearchManagerJni.get().getAcceptLanguages(
+                mNativeContextualSearchManagerPtr, this);
     }
 
     @Override
     public String getTranslateServiceTargetLanguage() {
-        return nativeGetTargetLanguage(mNativeContextualSearchManagerPtr);
+        return ContextualSearchManagerJni.get().getTargetLanguage(
+                mNativeContextualSearchManagerPtr, this);
     }
 
     // ============================================================================================
@@ -1036,8 +1039,9 @@ public class ContextualSearchManager
 
         @Override
         public void onContentViewCreated() {
-            nativeEnableContextualSearchJsApiForWebContents(
-                    mNativeContextualSearchManagerPtr, getSearchPanelWebContents());
+            ContextualSearchManagerJni.get().enableContextualSearchJsApiForWebContents(
+                    mNativeContextualSearchManagerPtr, ContextualSearchManager.this,
+                    getSearchPanelWebContents());
         }
 
         @Override
@@ -1575,8 +1579,9 @@ public class ContextualSearchManager
                 if (webContents != null) {
                     mInternalStateController.notifyStartingWorkOn(
                             InternalState.GATHERING_SURROUNDINGS);
-                    nativeGatherSurroundingText(
-                            mNativeContextualSearchManagerPtr, mContext, webContents);
+                    ContextualSearchManagerJni.get().gatherSurroundingText(
+                            mNativeContextualSearchManagerPtr, ContextualSearchManager.this,
+                            mContext, webContents);
                 } else {
                     mInternalStateController.reset(StateChangeReason.UNKNOWN);
                 }
@@ -1818,21 +1823,26 @@ public class ContextualSearchManager
         return mContext;
     }
 
-    // ============================================================================================
-    // Native calls
-    // ============================================================================================
-
-    private native long nativeInit();
-    private native void nativeDestroy(long nativeContextualSearchManager);
-    private native void nativeStartSearchTermResolutionRequest(long nativeContextualSearchManager,
-            ContextualSearchContext contextualSearchContext, WebContents baseWebContents);
-    protected native void nativeGatherSurroundingText(long nativeContextualSearchManager,
-            ContextualSearchContext contextualSearchContext, WebContents baseWebContents);
-    private native void nativeWhitelistContextualSearchJsApiUrl(
-            long nativeContextualSearchManager, String url);
-    private native void nativeEnableContextualSearchJsApiForWebContents(
-            long nativeContextualSearchManager, WebContents overlayWebContents);
-    // Don't call these directly, instead call the private methods that cache the results.
-    private native String nativeGetTargetLanguage(long nativeContextualSearchManager);
-    private native String nativeGetAcceptLanguages(long nativeContextualSearchManager);
+    @NativeMethods
+    interface Natives {
+        long init(ContextualSearchManager caller);
+
+        void destroy(long nativeContextualSearchManager, ContextualSearchManager caller);
+        void startSearchTermResolutionRequest(long nativeContextualSearchManager,
+                ContextualSearchManager caller, ContextualSearchContext contextualSearchContext,
+                WebContents baseWebContents);
+        void gatherSurroundingText(long nativeContextualSearchManager,
+                ContextualSearchManager caller, ContextualSearchContext contextualSearchContext,
+                WebContents baseWebContents);
+        void whitelistContextualSearchJsApiUrl(
+                long nativeContextualSearchManager, ContextualSearchManager caller, String url);
+        void enableContextualSearchJsApiForWebContents(long nativeContextualSearchManager,
+                ContextualSearchManager caller, WebContents overlayWebContents);
+        // Don't call these directly, instead call the private methods that cache the results.
+        String getTargetLanguage(
+                long nativeContextualSearchManager, ContextualSearchManager caller);
+
+        String getAcceptLanguages(
+                long nativeContextualSearchManager, ContextualSearchManager caller);
+    }
 }

chrome/android/javatests/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchTapEventTest.java

@@ -13,9 +13,12 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import org.chromium.base.test.util.CommandLineFlags;
 import org.chromium.base.test.util.Feature;
+import org.chromium.base.test.util.JniMocker;
 import org.chromium.base.test.util.Restriction;
 import org.chromium.chrome.browser.ChromeActivity;
 import org.chromium.chrome.browser.ChromeSwitches;
@@ -49,6 +52,12 @@ public class ContextualSearchTapEventTest {
     public ChromeActivityTestRule<ChromeActivity> mActivityTestRule =
             new ChromeActivityTestRule<>(ChromeActivity.class);
 
+    @Rule
+    public JniMocker mocker = new JniMocker();
+
+    @Mock
+    ContextualSearchManager.Natives mContextualSearchManagerJniMock;
+
     private ContextualSearchManagerWrapper mContextualSearchManager;
     private ContextualSearchPanel mPanel;
     private OverlayPanelManagerWrapper mPanelManager;
@@ -105,10 +114,6 @@ public class ContextualSearchTapEventTest {
                     "", "", "", "", QuickActionCategory.NONE, 0, "", "", 0);
         }
 
-        @Override
-        protected void nativeGatherSurroundingText(long nativeContextualSearchManager,
-                ContextualSearchContext contextualSearchContext, WebContents baseWebContents) {}
-
         /**
          * @return A stubbed SelectionPopupController for mocking text selection.
          */
@@ -247,6 +252,8 @@ public class ContextualSearchTapEventTest {
     public void setUp() throws Exception {
         mActivityTestRule.startMainActivityOnBlankPage();
         final ChromeActivity activity = mActivityTestRule.getActivity();
+        MockitoAnnotations.initMocks(this);
+        mocker.mock(ContextualSearchManagerJni.TEST_HOOKS, mContextualSearchManagerJniMock);
 
         TestThreadUtils.runOnUiThreadBlocking(() -> {
             mPanelManager = new OverlayPanelManagerWrapper();

chrome/android/junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchContextForTest.java

-// Copyright 2017 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.contextualsearch;
-
-/**
- * Provides a {@link ContextualSearchContext} suitable for injection into a class for testing.  This
- * has native calls stubbed out and is set up to provide some simple test feedback.
- */
-public class ContextualSearchContextForTest extends ContextualSearchContext {
-    private boolean mDidSelectionChange;
-    private String mDetectedLanguage;
-
-    /**
-     * @return Whether {@link #onSelectionChanged} was called.
-     */
-    boolean getDidSelectionChange() {
-        return mDidSelectionChange;
-    }
-
-    void setLanguageToDetect(String language) {
-        mDetectedLanguage = language;
-    }
-
-    @Override
-    void onSelectionChanged() {
-        mDidSelectionChange = true;
-    }
-
-    @Override
-    protected long nativeInit() {
-        return -1;
-    }
-
-    @Override
-    protected void nativeDestroy(long nativeContextualSearchContext) {}
-
-    @Override
-    protected void nativeSetResolveProperties(long nativeContextualSearchContext,
-            String homeCountry, boolean maySendBasePageUrl, long previousEventId,
-            int previousEventResults) {}
-
-    @Override
-    protected void nativeAdjustSelection(
-            long nativeContextualSearchContext, int startAdjust, int endAdjust) {}
-
-    @Override
-    protected String nativeDetectLanguage(long nativeContextualSearchContext) {
-        return mDetectedLanguage;
-    }
-}

chrome/android/junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchContextTest.java

@@ -9,13 +9,19 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.when;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.BlockJUnit4ClassRunner;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import org.chromium.base.test.util.Feature;
+import org.chromium.base.test.util.JniMocker;
 
 /**
  * Tests parts of the ContextualSearchContext class.
@@ -27,11 +33,30 @@ public class ContextualSearchContextTest {
     private static final String SAMPLE_TEXT =
             "Now Barack Obama is not the best example.  And Clinton is ambiguous.";
     private static final String HOME_COUNTRY = "unused";
+    private static final long NATIVE_PTR = 1;
 
-    private ContextualSearchContextForTest mContext;
+    private ContextualSearchContext mContext;
+    private boolean mDidSelectionChange;
+
+    private class ContextualSearchContextForTest extends ContextualSearchContext {
+        @Override
+        void onSelectionChanged() {
+            mDidSelectionChange = true;
+        }
+    }
+
+    @Rule
+    public JniMocker mocker = new JniMocker();
+
+    @Mock
+    private ContextualSearchContext.Natives mContextJniMock;
 
     @Before
     public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mocker.mock(ContextualSearchContextJni.TEST_HOOKS, mContextJniMock);
+        when(mContextJniMock.init(any())).thenReturn(NATIVE_PTR);
+        mDidSelectionChange = false;
         mContext = new ContextualSearchContextForTest();
     }
 
@@ -98,7 +123,7 @@ public class ContextualSearchContextTest {
         assertTrue(mContext.getSelectionStartOffset() >= 0);
         assertTrue(mContext.getSelectionEndOffset() >= 0);
         assertNotNull(mContext.getEncoding());
-        assertTrue(mContext.getDidSelectionChange());
+        assertTrue(mDidSelectionChange);
     }
 
     @Test
@@ -114,13 +139,13 @@ public class ContextualSearchContextTest {
         assertTrue(mContext.getSelectionEndOffset() >= 0);
         assertNotNull(mContext.getEncoding());
         assertNull(mContext.getInitialSelectedWord());
-        assertFalse(mContext.getDidSelectionChange());
+        assertFalse(mDidSelectionChange);
 
         simulateSelectWordAroundCaret(-"Ba".length(), "rack".length());
         assertEquals("Barack", mContext.getInitialSelectedWord());
         assertEquals("Barack".length(),
                 mContext.getSelectionEndOffset() - mContext.getSelectionStartOffset());
-        assertTrue(mContext.getDidSelectionChange());
+        assertTrue(mDidSelectionChange);
         assertTrue(mContext.hasValidSelection());
     }
 

chrome/android/junit/src/org/chromium/chrome/browser/contextualsearch/ContextualSearchEntityHeuristicTest.java

@@ -6,12 +6,20 @@ package org.chromium.chrome.browser.contextualsearch;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.when;
 
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.BlockJUnit4ClassRunner;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import org.chromium.base.test.util.Feature;
+import org.chromium.base.test.util.JniMocker;
 
 import java.util.Locale;
 
@@ -24,13 +32,26 @@ public class ContextualSearchEntityHeuristicTest {
             "Now Barack Obama, Michelle are not the best examples.  And Clinton is ambiguous.";
     private static final String UTF_8 = "UTF-8";
 
-    private ContextualSearchContextForTest mContext;
+    @Rule
+    public JniMocker mocker = new JniMocker();
+
+    @Mock
+    private ContextualSearchContext.Natives mContextJniMock;
+
+    private ContextualSearchContext mContext;
     private ContextualSearchEntityHeuristic mEntityHeuristic;
 
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mocker.mock(ContextualSearchContextJni.TEST_HOOKS, mContextJniMock);
+    }
+
     private void setupInstanceToTest(Locale locale, int tapOffset) {
-        mContext = new ContextualSearchContextForTest();
+        mContext = new ContextualSearchContext.ChangeIgnoringContext();
         mContext.setSurroundingText(UTF_8, SAMPLE_TEXT, tapOffset, tapOffset);
-        mContext.setLanguageToDetect(locale.getLanguage());
+        when(mContextJniMock.detectLanguage(anyLong(), eq(mContext)))
+                .thenReturn(locale.getLanguage());
         mEntityHeuristic = ContextualSearchEntityHeuristic.testInstance(mContext, true);
     }
 
@@ -88,8 +109,8 @@ public class ContextualSearchEntityHeuristicTest {
 
     private ContextualSearchEntityHeuristic setupHeuristic(
             String language, String start, String text) {
-        ContextualSearchContextForTest context = new ContextualSearchContextForTest();
-        context.setLanguageToDetect(language);
+        ContextualSearchContext context = new ContextualSearchContext.ChangeIgnoringContext();
+        when(mContextJniMock.detectLanguage(anyLong(), eq(context))).thenReturn(language);
         assert text.startsWith(start);
         int tapOffset = start.length();
         context.setSurroundingText(UTF_8, text, tapOffset, tapOffset);