python mock side_effect oder return_value abhängig von call_count

Question

Um eine Polling-Funktion zu testen, möchte ich den Aufruf einer Unterfunktion verspotten, so dass sie beim ersten Aufruf fehlschlägt und beim zweiten Aufruf erfolgreich ist. Hier ist eine sehr vereinfachte Version davon:

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 

Ist dies möglich, mit einem Mock Objekt aus dem mock Rahmen zu tun? Ich weiß Mock Objekte haben ein call_count Attribut, das ich irgendwie verwenden könnte.

Gerade jetzt habe ich es durch das Erstellen eines benutzerdefinierten Mock-Objekt, das ich Affe Patch verwenden gelöst sub_function(), aber ich fühle es eine bessere sein sollte weniger ausführlich Art und Weise tun:

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

Wenn ich verstehe, Ihre Frage richtig, Sie tun es durch setting side_effect to an iterable. Für Ihren einfachen Fall:

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

Wenn Sie für eine unbegrenzte Anzahl von Anrufen zulassen möchten, verwenden Sie die itertoolscycle und chain Funktionen:

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