[Autofill Assistant] Refactor |GetElementRect| to take Element

This refactors |WebController::GetElementRect| to take an
|ElementFinder::Result| instead of a |Selector|.

Bug: b/168107066
Change-Id: I0cbe7b12c95a92dac14c992017e41e3b0301f5e0
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2505650
Commit-Queue: Sandro Maggi <sandromaggi@google.com>
Reviewed-by: Clemens Arbesser <arbesser@google.com>
Cr-Commit-Position: refs/heads/master@{#822637}

components/autofill_assistant/browser/actions/action_delegate_util.cc

@@ -72,34 +72,6 @@ void FindElementAndPerformImpl(
       base::BindOnce(&OnFindElement, std::move(perform), std::move(done)));
 }
 
-template <typename R>
-void RetainElementAndExecuteGetCallback(
-    std::unique_ptr<ElementFinder::Result> element,
-    base::OnceCallback<void(const ClientStatus&, const R&)> callback,
-    const ClientStatus& status,
-    const R& result) {
-  DCHECK(element != nullptr);
-  std::move(callback).Run(status, result);
-}
-
-template <typename R>
-void TakeElementAndGetPropertyImpl(
-    ElementActionGetCallback<R> perform_and_get,
-    base::OnceCallback<void(const ClientStatus&, const R&)> done,
-    const ClientStatus& element_status,
-    std::unique_ptr<ElementFinder::Result> element_result) {
-  if (!element_status.ok()) {
-    VLOG(1) << __func__ << " Failed to find element.";
-    std::move(done).Run(element_status, R());
-    return;
-  }
-
-  std::move(perform_and_get)
-      .Run(*element_result,
-           base::BindOnce(&RetainElementAndExecuteGetCallback<R>,
-                          std::move(element_result), std::move(done)));
-}
-
 }  // namespace
 
 void PerformAll(std::unique_ptr<ElementActionVector> perform_actions,
@@ -125,16 +97,6 @@ void TakeElementAndPerform(ElementActionCallback perform,
                 std::move(element));
 }
 
-void TakeElementAndGetProperty(
-    ElementActionGetCallback<std::string> perform_and_get,
-    base::OnceCallback<void(const ClientStatus&, const std::string&)> done,
-    const ClientStatus& element_status,
-    std::unique_ptr<ElementFinder::Result> element) {
-  TakeElementAndGetPropertyImpl<std::string>(std::move(perform_and_get),
-                                             std::move(done), element_status,
-                                             std::move(element));
-}
-
 void ClickOrTapElement(const ActionDelegate* delegate,
                        const Selector& selector,
                        ClickType click_type,

components/autofill_assistant/browser/actions/action_delegate_util.h

@@ -12,6 +12,19 @@
 
 namespace autofill_assistant {
 namespace action_delegate_util {
+namespace {
+
+template <typename R>
+void RetainElementAndExecuteGetCallback(
+    std::unique_ptr<ElementFinder::Result> element,
+    base::OnceCallback<void(const ClientStatus&, const R&)> callback,
+    const ClientStatus& status,
+    const R& result) {
+  DCHECK(element != nullptr);
+  std::move(callback).Run(status, result);
+}
+
+}  // namespace
 
 using ElementActionCallback =
     base::OnceCallback<void(const ElementFinder::Result&,
@@ -45,11 +58,23 @@ void TakeElementAndPerform(ElementActionCallback perform,
 // element and the |done| callback as arguments, while retaining the element.
 // If the initial status is not ok, execute the |done| callback with the default
 // value immediately.
+template <typename R>
 void TakeElementAndGetProperty(
-    ElementActionGetCallback<std::string> perform_and_get,
-    base::OnceCallback<void(const ClientStatus&, const std::string&)> done,
+    ElementActionGetCallback<R> perform_and_get,
+    base::OnceCallback<void(const ClientStatus&, const R&)> done,
     const ClientStatus& element_status,
-    std::unique_ptr<ElementFinder::Result> element);
+    std::unique_ptr<ElementFinder::Result> element_result) {
+  if (!element_status.ok()) {
+    VLOG(1) << __func__ << " Failed to find element.";
+    std::move(done).Run(element_status, R());
+    return;
+  }
+
+  std::move(perform_and_get)
+      .Run(*element_result,
+           base::BindOnce(&RetainElementAndExecuteGetCallback<R>,
+                          std::move(element_result), std::move(done)));
+}
 
 void PerformAll(std::unique_ptr<ElementActionVector> perform_actions,
                 const ElementFinder::Result& element,

components/autofill_assistant/browser/actions/action_delegate_util_unittest.cc

@@ -215,7 +215,7 @@ TEST_F(ActionDelegateUtilTest, TakeElementAndGetProperty) {
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), "value"));
   EXPECT_CALL(*this, MockDoneGet(EqualsStatus(OkClientStatus()), "value"));
 
-  TakeElementAndGetProperty(
+  TakeElementAndGetProperty<std::string>(
       base::BindOnce(&ActionDelegateUtilTest::MockGetAction,
                      base::Unretained(this)),
       base::BindOnce(&ActionDelegateUtilTest::MockDoneGet,
@@ -233,7 +233,7 @@ TEST_F(ActionDelegateUtilTest, TakeElementAndGetPropertyWithFailedStatus) {
               MockDoneGet(EqualsStatus(ClientStatus(ELEMENT_RESOLUTION_FAILED)),
                           std::string()));
 
-  TakeElementAndGetProperty(
+  TakeElementAndGetProperty<std::string>(
       base::BindOnce(&ActionDelegateUtilTest::MockGetAction,
                      base::Unretained(this)),
       base::BindOnce(&ActionDelegateUtilTest::MockDoneGet,

components/autofill_assistant/browser/actions/get_element_status_action.cc

@@ -131,7 +131,7 @@ void GetElementStatusAction::OnWaitForElement(
   delegate_->FindElement(
       selector_,
       base::BindOnce(
-          &action_delegate_util::TakeElementAndGetProperty,
+          &action_delegate_util::TakeElementAndGetProperty<std::string>,
           base::BindOnce(&ActionDelegate::GetStringAttribute,
                          delegate_->GetWeakPtr(), attribute_list),
           base::BindOnce(&GetElementStatusAction::OnGetStringAttribute,

components/autofill_assistant/browser/actions/upload_dom_action.cc

@@ -47,11 +47,13 @@ void UploadDomAction::OnWaitForElement(const Selector& selector,
   }
 
   delegate_->FindElement(
-      selector, base::BindOnce(&action_delegate_util::TakeElementAndGetProperty,
-                               base::BindOnce(&ActionDelegate::GetOuterHtml,
-                                              delegate_->GetWeakPtr()),
-                               base::BindOnce(&UploadDomAction::OnGetOuterHtml,
-                                              weak_ptr_factory_.GetWeakPtr())));
+      selector,
+      base::BindOnce(
+          &action_delegate_util::TakeElementAndGetProperty<std::string>,
+          base::BindOnce(&ActionDelegate::GetOuterHtml,
+                         delegate_->GetWeakPtr()),
+          base::BindOnce(&UploadDomAction::OnGetOuterHtml,
+                         weak_ptr_factory_.GetWeakPtr())));
 }
 
 void UploadDomAction::OnGetOuterHtml(const ClientStatus& status,

components/autofill_assistant/browser/batch_element_checker.cc

@@ -65,15 +65,16 @@ void BatchElementChecker::Run(WebController* web_controller) {
   for (auto& entry : get_field_value_callbacks_) {
     web_controller->FindElement(
         entry.first, /* strict= */ true,
-        base::BindOnce(&action_delegate_util::TakeElementAndGetProperty,
-                       base::BindOnce(&WebController::GetFieldValue,
-                                      web_controller->GetWeakPtr()),
-                       base::BindOnce(&BatchElementChecker::OnFieldValueChecked,
-                                      weak_ptr_factory_.GetWeakPtr(),
-                                      // Guaranteed to exist for the lifetime of
-                                      // this instance, because the map isn't
-                                      // modified after Run has been called.
-                                      base::Unretained(&entry.second))));
+        base::BindOnce(
+            &action_delegate_util::TakeElementAndGetProperty<std::string>,
+            base::BindOnce(&WebController::GetFieldValue,
+                           web_controller->GetWeakPtr()),
+            base::BindOnce(&BatchElementChecker::OnFieldValueChecked,
+                           weak_ptr_factory_.GetWeakPtr(),
+                           // Guaranteed to exist for the lifetime of
+                           // this instance, because the map isn't
+                           // modified after Run has been called.
+                           base::Unretained(&entry.second))));
   }
 
   // The extra +1 of pending_check_count and this check happening last

components/autofill_assistant/browser/element_area.cc

@@ -12,6 +12,7 @@
 #include "base/single_thread_task_runner.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
+#include "components/autofill_assistant/browser/actions/action_delegate_util.h"
 #include "components/autofill_assistant/browser/script_executor_delegate.h"
 #include "components/autofill_assistant/browser/web/web_controller.h"
 
@@ -128,10 +129,15 @@ void ElementArea::Update() {
 
   for (auto& rectangle : rectangles_) {
     for (auto& position : rectangle.positions) {
-      delegate_->GetWebController()->GetElementRect(
-          position.selector,
-          base::BindOnce(&ElementArea::OnGetElementRect,
-                         weak_ptr_factory_.GetWeakPtr(), position.selector));
+      delegate_->GetWebController()->FindElement(
+          position.selector, /* strict= */ true,
+          base::BindOnce(
+              &action_delegate_util::TakeElementAndGetProperty<RectF>,
+              base::BindOnce(&WebController::GetElementRect,
+                             delegate_->GetWebController()->GetWeakPtr()),
+              base::BindOnce(&ElementArea::OnGetElementRect,
+                             weak_ptr_factory_.GetWeakPtr(),
+                             position.selector)));
     }
   }
 }

components/autofill_assistant/browser/element_area_unittest.cc

@@ -13,6 +13,7 @@
 #include "base/test/gmock_callback_support.h"
 #include "base/test/mock_callback.h"
 #include "base/test/task_environment.h"
+#include "components/autofill_assistant/browser/actions/action_test_utils.h"
 #include "components/autofill_assistant/browser/fake_script_executor_delegate.h"
 #include "components/autofill_assistant/browser/script_executor_delegate.h"
 #include "components/autofill_assistant/browser/web/mock_web_controller.h"
@@ -70,6 +71,7 @@ class ElementAreaTest : public testing::Test {
     delegate_.GetMutableSettings()->element_position_update_interval =
         base::TimeDelta::FromMilliseconds(100);
 
+    test_util::MockFindAnyElement(mock_web_controller_);
     ON_CALL(mock_web_controller_, OnGetElementRect(_, _))
         .WillByDefault(
             RunOnceCallback<1>(ClientStatus(UNEXPECTED_JS_ERROR), RectF()));
@@ -145,8 +147,13 @@ TEST_F(ElementAreaTest, GetVisualViewport) {
 }
 
 TEST_F(ElementAreaTest, OneRectangle) {
+  Selector expected_selector({"#found"});
+  expected_selector.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#found"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(25, 25, 75, 75)));
 
   SetElement("#found");
@@ -156,8 +163,13 @@ TEST_F(ElementAreaTest, OneRectangle) {
 }
 
 TEST_F(ElementAreaTest, CallOnUpdate) {
+  Selector expected_selector({"#found"});
+  expected_selector.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#found"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(25, 25, 75, 75)));
 
   SetElement("#found");
@@ -167,8 +179,14 @@ TEST_F(ElementAreaTest, CallOnUpdate) {
 }
 
 TEST_F(ElementAreaTest, CallOnUpdateAfterSetFromProto) {
+  Selector expected_selector({"#found"});
+  expected_selector.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#found"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector, 2)),
+                               _))
+      .Times(2)
       .WillRepeatedly(
           RunOnceCallback<1>(OkClientStatus(), RectF(25, 25, 75, 75)));
 
@@ -179,12 +197,17 @@ TEST_F(ElementAreaTest, CallOnUpdateAfterSetFromProto) {
 }
 
 TEST_F(ElementAreaTest, DontCallOnUpdateWhenViewportMissing) {
+  Selector expected_selector({"#found"});
+  expected_selector.MustBeVisible();
+
   // Swallowing calls to OnGetVisualViewport guarantees that the viewport
   // position will never be known.
   EXPECT_CALL(mock_web_controller_, OnGetVisualViewport(_))
       .WillOnce(DoNothing());
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#found"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(25, 25, 75, 75)));
 
   SetElement("#found");
@@ -208,12 +231,22 @@ TEST_F(ElementAreaTest, CallOnUpdateWhenViewportMissingAndEmptyRect) {
 }
 
 TEST_F(ElementAreaTest, TwoRectangles) {
-  EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#top_left"}).MustBeVisible()), _))
-      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 0, 25, 25)));
+  Selector expected_selector_top_left({"#top_left"});
+  expected_selector_top_left.MustBeVisible();
+  Selector expected_selector_bottom_right({"#bottom_right"});
+  expected_selector_bottom_right.MustBeVisible();
+
   EXPECT_CALL(
       mock_web_controller_,
-      OnGetElementRect(Eq(Selector({"#bottom_right"}).MustBeVisible()), _))
+      OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                           mock_web_controller_, expected_selector_top_left)),
+                       _))
+      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 0, 25, 25)));
+  EXPECT_CALL(mock_web_controller_,
+              OnGetElementRect(
+                  EqualsElement(test_util::MockFindElement(
+                      mock_web_controller_, expected_selector_bottom_right)),
+                  _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(25, 25, 100, 100)));
 
   ElementAreaProto area_proto;
@@ -229,11 +262,20 @@ TEST_F(ElementAreaTest, TwoRectangles) {
 }
 
 TEST_F(ElementAreaTest, OneRectangleTwoElements) {
+  Selector expected_selector_1({"#element1"});
+  expected_selector_1.MustBeVisible();
+  Selector expected_selector_2({"#element2"});
+  expected_selector_2.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element1"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_1)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(1, 3, 2, 4)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element2"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_2)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(5, 2, 6, 5)));
 
   ElementAreaProto area_proto;
@@ -248,14 +290,22 @@ TEST_F(ElementAreaTest, OneRectangleTwoElements) {
 }
 
 TEST_F(ElementAreaTest, DoNotReportIncompleteRectangles) {
+  Selector expected_selector_1({"#element1"});
+  expected_selector_1.MustBeVisible();
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element1"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_1)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(1, 3, 2, 4)));
 
   // Getting the position of #element2 neither succeeds nor fails, simulating an
   // intermediate state which shouldn't be reported to the callback.
+  Selector expected_selector_2({"#element2"});
+  expected_selector_2.MustBeVisible();
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element2"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_2)),
+                               _))
       .WillOnce(DoNothing());  // overrides default action
 
   ElementAreaProto area_proto;
@@ -272,17 +322,34 @@ TEST_F(ElementAreaTest, DoNotReportIncompleteRectangles) {
 }
 
 TEST_F(ElementAreaTest, OneRectangleFourElements) {
+  Selector expected_selector_1({"#element1"});
+  expected_selector_1.MustBeVisible();
+  Selector expected_selector_2({"#element2"});
+  expected_selector_2.MustBeVisible();
+  Selector expected_selector_3({"#element3"});
+  expected_selector_3.MustBeVisible();
+  Selector expected_selector_4({"#element4"});
+  expected_selector_4.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element1"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_1)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 0, 1, 1)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element2"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_2)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(9, 9, 100, 100)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element3"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_3)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 9, 1, 100)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element4"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_4)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(9, 0, 100, 1)));
 
   ElementAreaProto area_proto;
@@ -299,11 +366,20 @@ TEST_F(ElementAreaTest, OneRectangleFourElements) {
 }
 
 TEST_F(ElementAreaTest, OneRectangleMissingElementsReported) {
+  Selector expected_selector_1({"#element1"});
+  expected_selector_1.MustBeVisible();
+  Selector expected_selector_2({"#element2"});
+  expected_selector_2.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element1"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_1)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(1, 1, 2, 2)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element2"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_2)),
+                               _))
       .WillOnce(RunOnceCallback<1>(ClientStatus(UNEXPECTED_JS_ERROR), RectF()));
 
   ElementAreaProto area_proto;
@@ -320,11 +396,20 @@ TEST_F(ElementAreaTest, OneRectangleMissingElementsReported) {
 }
 
 TEST_F(ElementAreaTest, FullWidthRectangle) {
+  Selector expected_selector_1({"#element1"});
+  expected_selector_1.MustBeVisible();
+  Selector expected_selector_2({"#element2"});
+  expected_selector_2.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element1"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_1)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(1, 3, 2, 4)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element2"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector_2)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(5, 7, 6, 8)));
   EXPECT_CALL(mock_web_controller_, OnGetVisualViewport(_))
       .WillRepeatedly(
@@ -347,9 +432,19 @@ TEST_F(ElementAreaTest, FullWidthRectangle) {
 
 TEST_F(ElementAreaTest, ElementMovesAfterUpdate) {
   testing::InSequence seq;
+
+  Selector expected_selector({"#element"});
+  expected_selector.MustBeVisible();
+
+  EXPECT_CALL(mock_web_controller_,
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
+      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 25, 100, 50)));
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element"}).MustBeVisible()), _))
-      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 25, 100, 50)))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 50, 100, 75)));
 
   SetElement("#element");
@@ -372,11 +467,25 @@ TEST_F(ElementAreaTest, ElementMovesAfterUpdate) {
 
 TEST_F(ElementAreaTest, ElementMovesWithTime) {
   testing::InSequence seq;
+
+  Selector expected_selector({"#element"});
+  expected_selector.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(Eq(Selector({"#element"}).MustBeVisible()), _))
-      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 25, 100, 50)))
-      .WillRepeatedly(
-          RunOnceCallback<1>(OkClientStatus(), RectF(0, 50, 100, 75)));
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
+      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 25, 100, 50)));
+  EXPECT_CALL(mock_web_controller_,
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
+      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 50, 100, 75)));
+  EXPECT_CALL(mock_web_controller_,
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
+      .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(0, 50, 100, 75)));
 
   SetElement("#element");
 
@@ -399,9 +508,13 @@ TEST_F(ElementAreaTest, ElementMovesWithTime) {
 }
 
 TEST_F(ElementAreaTest, RestrictedElement) {
+  Selector expected_selector({"#restricted_element"});
+  expected_selector.MustBeVisible();
+
   EXPECT_CALL(mock_web_controller_,
-              OnGetElementRect(
-                  Eq(Selector({"#restricted_element"}).MustBeVisible()), _))
+              OnGetElementRect(EqualsElement(test_util::MockFindElement(
+                                   mock_web_controller_, expected_selector)),
+                               _))
       .WillOnce(RunOnceCallback<1>(OkClientStatus(), RectF(25, 25, 75, 75)));
 
   SetElement("#restricted_element", /* restricted= */ true);

components/autofill_assistant/browser/web/mock_web_controller.h

@@ -107,12 +107,12 @@ class MockWebController : public WebController {
                         callback));
 
   void GetElementRect(
-      const Selector& selector,
+      const ElementFinder::Result& element,
       ElementRectGetter::ElementRectCallback callback) override {
-    OnGetElementRect(selector, callback);
+    OnGetElementRect(element, callback);
   }
   MOCK_METHOD2(OnGetElementRect,
-               void(const Selector& selector,
+               void(const ElementFinder::Result& element,
                     ElementRectGetter::ElementRectCallback& callback));
 
   void WaitForWindowHeightChange(

components/autofill_assistant/browser/web/web_controller.cc

@@ -1317,28 +1317,15 @@ void WebController::OnGetVisualViewport(
 }
 
 void WebController::GetElementRect(
-    const Selector& selector,
+    const ElementFinder::Result& element,
     ElementRectGetter::ElementRectCallback callback) {
-  FindElement(
-      selector, /* strict_mode= */ true,
-      base::BindOnce(&WebController::OnFindElementForRect,
-                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
-}
-
-void WebController::OnFindElementForRect(
-    ElementRectGetter::ElementRectCallback callback,
-    const ClientStatus& element_status,
-    std::unique_ptr<ElementFinder::Result> element_result) {
-  if (!element_status.ok()) {
-    std::move(callback).Run(element_status, RectF());
-    return;
-  }
   std::unique_ptr<ElementRectGetter> getter =
       std::make_unique<ElementRectGetter>(devtools_client_.get());
   auto* ptr = getter.get();
   pending_workers_.emplace_back(std::move(getter));
   ptr->Start(
-      std::move(element_result),
+      // TODO(b/172041811): Ownership of element.
+      std::make_unique<ElementFinder::Result>(element),
       base::BindOnce(&WebController::OnGetElementRect,
                      weak_ptr_factory_.GetWeakPtr(), ptr, std::move(callback)));
 }

components/autofill_assistant/browser/web/web_controller.h

@@ -215,14 +215,14 @@ class WebController {
   virtual void GetVisualViewport(
       base::OnceCallback<void(const ClientStatus&, const RectF&)> callback);
 
-  // Gets the position of the element identified by the selector.
+  // Gets the position of the |element|.
   //
   // If unsuccessful, the callback gets the failure status with an empty rect.
   //
   // If successful, the callback gets a success status with a set of
   // (left, top, right, bottom) coordinates rect, expressed in absolute CSS
   // coordinates.
-  virtual void GetElementRect(const Selector& selector,
+  virtual void GetElementRect(const ElementFinder::Result& element,
                               ElementRectGetter::ElementRectCallback callback);
 
   // Calls the callback once the main document window has been resized.
@@ -378,9 +378,6 @@ class WebController {
       size_t index,
       int delay_in_milli,
       base::OnceCallback<void(const ClientStatus&)> callback);
-  void OnFindElementForRect(ElementRectGetter::ElementRectCallback callback,
-                            const ClientStatus& status,
-                            std::unique_ptr<ElementFinder::Result> result);
   void OnGetElementRect(ElementRectGetter* getter_to_release,
                         ElementRectGetter::ElementRectCallback callback,
                         const ClientStatus& rect_status,

components/autofill_assistant/browser/web/web_controller_browsertest.cc

@@ -800,15 +800,39 @@ class WebControllerBrowserTest : public content::ContentBrowserTest,
   ClientStatus GetElementRect(const Selector& selector, RectF* rect_output) {
     base::RunLoop run_loop;
     ClientStatus result;
-    web_controller_->GetElementRect(
-        selector, base::BindOnce(&WebControllerBrowserTest::OnGetElementRect,
-                                 base::Unretained(this), run_loop.QuitClosure(),
-                                 &result, rect_output));
+
+    web_controller_->FindElement(
+        selector, /* strict= */ true,
+        base::BindOnce(&WebControllerBrowserTest::GetElementRectElementCallback,
+                       base::Unretained(this), run_loop.QuitClosure(), &result,
+                       rect_output));
+
     run_loop.Run();
     return result;
   }
 
-  void OnGetElementRect(base::OnceClosure done_callback,
+  void GetElementRectElementCallback(
+      base::OnceClosure done_callback,
+      ClientStatus* result_output,
+      RectF* rect_output,
+      const ClientStatus& element_status,
+      std::unique_ptr<ElementFinder::Result> element_result) {
+    if (!element_status.ok()) {
+      *result_output = element_status;
+      std::move(done_callback).Run();
+      return;
+    }
+
+    ASSERT_TRUE(element_result != nullptr);
+    web_controller_->GetElementRect(
+        *element_result,
+        base::BindOnce(&WebControllerBrowserTest::OnGetElementRect,
+                       base::Unretained(this), std::move(element_result),
+                       std::move(done_callback), result_output, rect_output));
+  }
+
+  void OnGetElementRect(std::unique_ptr<ElementFinder::Result> element,
+                        base::OnceClosure done_callback,
                         ClientStatus* result_output,
                         RectF* rect_output,
                         const ClientStatus& rect_status,