python mock side_effect o return_value dipendente dal call_count

Question

Per testare una funzione di polling, voglio prendere in giro il richiamo di una funzione secondaria in modo che la prima volta che viene chiamato non riesca e la seconda volta che viene chiamato, avrà esito positivo. Ecco una versione molto semplificata di esso:

poll_function(var1): 
    value = sub_function(var1) # First call will return None 
    while not value: 
     time.sleep(POLLING_INTERVAL) 
     value = sub_function(var1) # A subsequent call will return a string, e.g "data" 
    return value 

Questo è possibile a che fare con un oggetto Mock dal quadro mock? So che gli oggetti Mock hanno un attributo call_count che dovrei essere in grado di usare in qualche modo.

In questo momento ho risolto con la creazione di un oggetto fittizio personalizzato che uso per cerotto scimmia sub_function(), ma mi sento ci dovrebbe essere un modo migliore meno prolissa di farlo:

def test_poll(): 
    class MyMock(object):              

     def __init__(self, *args):            
      self.call_count = 0             

     def sub_function(self, *args, **kwargs):        
      if self.call_count > 1:            
       return "data"    
      else:                
       self.call_count += 1           
       return None 

    my_mock = MyMock()               
    with patch('sub_function', my_mock.sub_function):   
     ok_(poll_function())   

Answer 1

Se ho capito la tua domanda correttamente, lo fai da setting side_effect to an iterable. Per il vostro semplice caso:

>>> mock_poll = Mock(side_effect=[None, 'data']) 
>>> mock_poll() 
None 
>>> mock_poll() 
'data' 

Se si desidera consentire per un numero illimitato di chiamate, utilizzare i itertoolscycle e chain funzioni:

>>> mock_poll = Mock(side_effect=chain(['first'], cycle(['others'])))