Fix flaky test under Python 3.12.

The functionality to enable networking (i.e. turn off airplane mode and turn on
wifi) is done asynchronously within `A11yGrpcWrapper._fetch_task_extras()`.
Under some conditions, repeated calls could interfere with each other before
the `concurrent.futures.Future` was reset, and this caused one test case to
fail.

This change makes this 100% robust by moving the reset to the top of the
function to avoid having the reset trigger within the same call of
`_fetch_task_extras()`.

PiperOrigin-RevId: 716214824

android_env/wrappers/a11y_grpc_wrapper.py

@@ -222,7 +222,7 @@ def _enable_a11y_tree_logs(self) -> None:
       )
 
   def _reset_enable_networking_attempts(self) -> None:
-    self._enable_networking_attepts_left = self._max_enable_networking_attempts
+    self._enable_networking_attempts_left = self._max_enable_networking_attempts
     self._enabling_networking_future = None
     self._a11y_exception = None
 
@@ -446,6 +446,15 @@ def _fetch_task_extras(self) -> dict[str, Any]:
         services by re-enabling the network connection.
     """
     base_extras = self._env.task_extras(latest_only=False).copy()
+    # If the previous future is done, reset it to the initial state.
+    if (
+        self._enabling_networking_future is not None
+        and self._enabling_networking_future.done()
+    ):
+      self._enabling_networking_future = None
+      self._enable_networking_attempts_left -= 1
+      logging.info('Finished enabling networking.')
+
     if (
         self._enabling_networking_future is None
         and 'exception' in base_extras
@@ -455,10 +464,10 @@ def _fetch_task_extras(self) -> dict[str, Any]:
       logging.warning(
           'AccessibilityForwarder logged exceptions: %s', self._a11y_exception
       )
-      if self._enable_networking_attepts_left > 0:
+      if self._enable_networking_attempts_left > 0:
         logging.warning(
-            'Attempting to enable networking. %s attemps left.',
-            self._enable_networking_attepts_left - 1,
+            'Attempting to enable networking. %s attempts left.',
+            self._enable_networking_attempts_left - 1,
         )
         executor = futures.ThreadPoolExecutor(max_workers=1)
         self._enabling_networking_future = executor.submit(
@@ -470,14 +479,6 @@ def _fetch_task_extras(self) -> dict[str, Any]:
             f' exception.{self._a11y_exception}.'
         )
 
-    if (
-        self._enabling_networking_future is not None
-        and self._enabling_networking_future.done()
-    ):
-      self._enabling_networking_future = None
-      self._enable_networking_attepts_left -= 1
-      logging.info('Finished enabling networking.')
-
     forests = self._servicer.gather_forests()
     if forests:
       base_extras.update(a11y_forests.package_forests_to_task_extras(forests))

android_env/wrappers/a11y_grpc_wrapper_test.py

@@ -327,11 +327,8 @@ def test_fetch_task_extras_enable_networking_twice(
       a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True
   )
   @mock.patch.object(grpc, 'server', autospec=True)
-  def test_task_extras_rasises_with_a11y_info_exception(
-      self,
-      mock_sleep,
-      mock_add_servicer,
-      mock_server,
+  def test_task_extras_raises_a11y_info_exception(
+      self, mock_sleep, mock_add_servicer, mock_server
   ):
     del mock_server, mock_add_servicer, mock_sleep
     base_env = mock.create_autospec(
@@ -366,6 +363,7 @@ def test_task_extras_rasises_with_a11y_info_exception(
       extras = wrapped_env._fetch_task_extras()
       self.assertNotIn('accessibility_tree', extras)
       self.assertNotIn('full_event', extras)
+      # Wait for the the attempt to finish.
       future = wrapped_env._enabling_networking_future
       if future is not None:
         future.result()