code: async def check_and_add_pump_async(self, partition_id, lease):
          
          if partition_id in self.partition_pumps:
              # There already is a pump. Make sure the pump is working and replace the lease.
              captured_pump = self.partition_pumps[partition_id]
              if captured_pump.pump_status == "Errored" or captured_pump.is_closing():
                  # The existing pump is bad. Remove it.
                  await self.remove_pump_async(partition_id, "Shutdown")
              else:
                  # Pump is working, should just replace the lease.
                  # This is causing a race condition since if the checkpoint is being updated
                  # when the lease changes then the pump will error and shut down
                  captured_pump.set_lease(lease)
          else:
              _logger.info("Starting pump %r %r", self.host.guid, partition_id)
              await self.create_new_pump_async(partition_id, lease)
docs: Updates the lease on an exisiting pump.
  
          :param partition_id: The partition ID.
          :type partition_id: str
          :param lease: The lease to be used.
          :type lease: ~azure.eventprocessorhost.lease.Lease