[Payments] Use the correct index when checking addresses

chrome/android/javatests/src/org/chromium/chrome/browser/payments/PaymentRequestTestBase.java

Index: chrome/android/javatests/src/org/chromium/chrome/browser/payments/PaymentRequestTestBase.java
diff --git a/chrome/android/javatests/src/org/chromium/chrome/browser/payments/PaymentRequestTestBase.java b/chrome/android/javatests/src/org/chromium/chrome/browser/payments/PaymentRequestTestBase.java
index a138fb42285c69fb0b5e54d514be78f3174df255..243160e525201e977902e51d5ffa3ea1f26c6a94 100644
--- a/chrome/android/javatests/src/org/chromium/chrome/browser/payments/PaymentRequestTestBase.java
+++ b/chrome/android/javatests/src/org/chromium/chrome/browser/payments/PaymentRequestTestBase.java
@@ -20,6 +20,7 @@ import org.chromium.chrome.browser.autofill.CardUnmaskPrompt.CardUnmaskObserverF
 import org.chromium.chrome.browser.payments.PaymentRequestImpl.PaymentRequestServiceObserverForTest;
 import org.chromium.chrome.browser.payments.ui.EditorTextField;
 import org.chromium.chrome.browser.payments.ui.PaymentRequestSection.OptionSection;
+import org.chromium.chrome.browser.payments.ui.PaymentRequestSection.OptionSection.OptionRow;
 import org.chromium.chrome.browser.payments.ui.PaymentRequestUI;
 import org.chromium.chrome.browser.payments.ui.PaymentRequestUI.PaymentRequestObserverForTest;
 import org.chromium.chrome.test.ChromeActivityTestCaseBase;
@@ -45,6 +46,7 @@ abstract class PaymentRequestTestBase extends ChromeActivityTestCaseBase<ChromeA
     protected final PaymentsCallbackHelper<PaymentRequestUI> mReadyForInput;
     protected final PaymentsCallbackHelper<PaymentRequestUI> mReadyToPay;
     protected final PaymentsCallbackHelper<PaymentRequestUI> mReadyToClose;
+    protected final PaymentsCallbackHelper<PaymentRequestUI> mSelectionChecked;
     protected final PaymentsCallbackHelper<PaymentRequestUI> mResultReady;
     protected final PaymentsCallbackHelper<CardUnmaskPrompt> mReadyForUnmaskInput;
     protected final PaymentsCallbackHelper<CardUnmaskPrompt> mReadyToUnmask;
@@ -67,6 +69,7 @@ abstract class PaymentRequestTestBase extends ChromeActivityTestCaseBase<ChromeA
         mReadyForInput = new PaymentsCallbackHelper<>();
         mReadyToPay = new PaymentsCallbackHelper<>();
         mReadyToClose = new PaymentsCallbackHelper<>();
+        mSelectionChecked = new PaymentsCallbackHelper<>();
         mResultReady = new PaymentsCallbackHelper<>();
         mReadyForUnmaskInput = new PaymentsCallbackHelper<>();
         mReadyToUnmask = new PaymentsCallbackHelper<>();
@@ -288,7 +291,28 @@ abstract class PaymentRequestTestBase extends ChromeActivityTestCaseBase<ChromeA
     }
 
     /**
-     *  Returns the the number of shipping address suggestions,
+     *  Clicks on the label corresponding to the shipping address suggestion at the specified
+     *  |suggestionIndex|.
+     * @throws InterruptedException
+     */
+    protected void clickOnShippingAddressSuggestionOptionAndWait(
+            final int suggestionIndex, CallbackHelper helper)
+                    throws ExecutionException, TimeoutException, InterruptedException {
+        assert (suggestionIndex < getNumberOfShippingAddressSuggestions());
+
+        int callCount = helper.getCallCount();
+        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+            @Override
+            public void run() {
+                ((OptionSection) mUI.getShippingAddressSectionForTest())
+                        .getOptionLabelsForTest(suggestionIndex).performClick();
+            }
+        });
+        helper.waitForCallback(callCount);
+    }
+
+    /**
+     *  Returns the the number of shipping address suggestions.
      */
     protected int getNumberOfShippingAddressSuggestions() throws ExecutionException {
         return ThreadUtils.runOnUiThreadBlocking(new Callable<Integer>() {
@@ -300,6 +324,18 @@ abstract class PaymentRequestTestBase extends ChromeActivityTestCaseBase<ChromeA
         });
     }
 
+    /** Returns the {@link OptionRow} at the given index for the shipping address section. */
+    protected OptionRow getShippingAddressOptionRowAtIndex(final int index)
+            throws ExecutionException {
+        return ThreadUtils.runOnUiThreadBlocking(new Callable<OptionRow>() {
+            @Override
+            public OptionRow call() {
+                return ((OptionSection) mUI.getShippingAddressSectionForTest())
+                        .getOptionRowAtIndex(index);
+            }
+        });
+    }
+
     /** Selects the spinner value in the editor UI for credit cards. */
     protected void setSpinnerSelectionsInCardEditorAndWait(final int[] selections,
             CallbackHelper helper) throws InterruptedException, TimeoutException {
@@ -460,6 +496,12 @@ abstract class PaymentRequestTestBase extends ChromeActivityTestCaseBase<ChromeA
     }
 
     @Override
+    public void onPaymentRequestSelectionChecked(PaymentRequestUI ui) {
+        ThreadUtils.assertOnUiThread();
+        mSelectionChecked.notifyCalled(ui);
+    }
+
+    @Override
     public void onPaymentRequestResultReady(PaymentRequestUI ui) {
         ThreadUtils.assertOnUiThread();
         mResultReady.notifyCalled(ui);