code: def test_retry():
      helper = RetrySamplingHelper(10)
      assert helper.retry(lambda: 1) == 1
  
      class _keep_trying_and_will_success:
          def __init__(self):
              self._count = 0
  
          def __call__(self):
              self._count += 1
              if self._count < 5:
                  raise KeyError()
              return 1
  
      with pytest.raises(KeyError):
          helper.retry(_keep_trying_and_will_success())
  
      helper = RetrySamplingHelper(10, KeyError)
      assert helper.retry(_keep_trying_and_will_success()) == 1
  
      helper = RetrySamplingHelper(3, KeyError)
      assert helper.retry(_keep_trying_and_will_success()) is None
  
      helper = RetrySamplingHelper(3, KeyError, raise_last=True)
      with pytest.raises(KeyError):
          helper.retry(_keep_trying_and_will_success())
  
imports: import pytest
  from nni.mutable import *
  from nni.nas.strategy.utils import *
label_desc: Use image node_id 2 for calling a global function with example usage: test_retry() and returns: int