
from selenium.common import exceptions

from lib import base
from lib.constants import locator


class ExtendedInfo(base.Component):
  """Model representing an extended info box that allows the object to be
  mapped"""
  _locator = locator.ExtendedInfo

  def __init__(self, driver):
    super(ExtendedInfo, self).__init__(driver)
    self.button_map = None

  def _reload_contents(self):
    self.button_map = base.Button(
        self._driver, self._locator.BUTTON_MAP_TO)

  def map_to_object(self):
    try:
      self.button_map = base.Button(
          self._driver, self._locator.BUTTON_MAP_TO)
      self.button_map.click()
    except exceptions.StaleElementReferenceException:
      self._reload_contents()
      return self.map_to_object()

  def is_already_mapped(self):
    """Checks if the object is already mapped"""
    try:
      self._driver.find_element(*self._locator.ALREADY_MAPPED)
      return True
    except exceptions.NoSuchElementException:
      return False


from selenium.common import exceptions

from lib import base
from lib.constants import locator
from lib.utils import selenium_utils


class ExtendedInfo(base.Component):
  """Model representing an extended info box that allows the object to be
  mapped"""
  locator_cls = locator.ExtendedInfo

  def __init__(self, driver):
    super(ExtendedInfo, self).__init__(driver)
    self.is_mapped = None
    self.button_map = None
    self.title = base.Label(driver, self.locator_cls.TITLE)

    self._set_is_mapped()

  def map_to_object(self):
    selenium_utils.click_on_staleable_element(
        self._driver,
        self.locator_cls.BUTTON_MAP_TO)
    self.is_mapped = True

  def _set_is_mapped(self):
    """Checks if the object is already mapped"""
    try:
      self._driver.find_element(*self.locator_cls.ALREADY_MAPPED)
      self.is_mapped = True
    except exceptions.NoSuchElementException:
      self.is_mapped = False

Handle stealable element with utils
Handle stealable element with utils

  
  from selenium.common import exceptions
  
  from lib import base
  from lib.constants import locator
+ from lib.utils import selenium_utils
  
  
  class ExtendedInfo(base.Component):
    """Model representing an extended info box that allows the object to be
    mapped"""
-   _locator = locator.ExtendedInfo
+   locator_cls = locator.ExtendedInfo
  
    def __init__(self, driver):
      super(ExtendedInfo, self).__init__(driver)
+     self.is_mapped = None
      self.button_map = None
+     self.title = base.Label(driver, self.locator_cls.TITLE)
  
+     self._set_is_mapped()
-   def _reload_contents(self):
-     self.button_map = base.Button(
-         self._driver, self._locator.BUTTON_MAP_TO)
  
    def map_to_object(self):
-     try:
-       self.button_map = base.Button(
+     selenium_utils.click_on_staleable_element(
+         self._driver,
-           self._driver, self._locator.BUTTON_MAP_TO)
+         self.locator_cls.BUTTON_MAP_TO)
+     self.is_mapped = True
-       self.button_map.click()
-     except exceptions.StaleElementReferenceException:
-       self._reload_contents()
-       return self.map_to_object()
  
-   def is_already_mapped(self):
+   def _set_is_mapped(self):
      """Checks if the object is already mapped"""
      try:
-       self._driver.find_element(*self._locator.ALREADY_MAPPED)
+       self._driver.find_element(*self.locator_cls.ALREADY_MAPPED)
-       return True
+       self.is_mapped = True
      except exceptions.NoSuchElementException:
-       return False
+       self.is_mapped = False
  
Handle stealable element with utils
## Code Before:

from selenium.common import exceptions

from lib import base
from lib.constants import locator


class ExtendedInfo(base.Component):
  """Model representing an extended info box that allows the object to be
  mapped"""
  _locator = locator.ExtendedInfo

  def __init__(self, driver):
    super(ExtendedInfo, self).__init__(driver)
    self.button_map = None

  def _reload_contents(self):
    self.button_map = base.Button(
        self._driver, self._locator.BUTTON_MAP_TO)

  def map_to_object(self):
    try:
      self.button_map = base.Button(
          self._driver, self._locator.BUTTON_MAP_TO)
      self.button_map.click()
    except exceptions.StaleElementReferenceException:
      self._reload_contents()
      return self.map_to_object()

  def is_already_mapped(self):
    """Checks if the object is already mapped"""
    try:
      self._driver.find_element(*self._locator.ALREADY_MAPPED)
      return True
    except exceptions.NoSuchElementException:
      return False

## Instruction:
Handle stealable element with utils
## Code After:

from selenium.common import exceptions

from lib import base
from lib.constants import locator
from lib.utils import selenium_utils


class ExtendedInfo(base.Component):
  """Model representing an extended info box that allows the object to be
  mapped"""
  locator_cls = locator.ExtendedInfo

  def __init__(self, driver):
    super(ExtendedInfo, self).__init__(driver)
    self.is_mapped = None
    self.button_map = None
    self.title = base.Label(driver, self.locator_cls.TITLE)

    self._set_is_mapped()

  def map_to_object(self):
    selenium_utils.click_on_staleable_element(
        self._driver,
        self.locator_cls.BUTTON_MAP_TO)
    self.is_mapped = True

  def _set_is_mapped(self):
    """Checks if the object is already mapped"""
    try:
      self._driver.find_element(*self.locator_cls.ALREADY_MAPPED)
      self.is_mapped = True
    except exceptions.NoSuchElementException:
      self.is_mapped = False
