Comment faire pour que MagicMock de Python renvoie une valeur la première fois qu'il est appelé et déclenche une exception la deuxième fois qu'il est appelé?

Question

J'ai le programme python suivant qui utilise le mocking.

#!/usr/bin/env python 
import mock 

def my_func1(): 
    return "Hello" 

my_func = mock.MagicMock() 
my_func.return_value = "Goodbye" 

print my_func() 
print my_func() 

Sortie:

Goodbye 
Goodbye 

tout fonctionne comme il se doit. Génial.

Mais je veux que la méthode simulée renvoie Goodbye la première fois qu'elle est appelée et déclenche une exception la deuxième fois qu'elle est appelée. Comment puis je faire ça??

Answer 1

Comme Sraw a souligné, vous pouvez utiliser side_effect. Je serais probablement utiliser une fonction de générateur plutôt que d'introduire une approche globale:

import mock 

def effect(*args, **kwargs): 
    yield "Goodbye" 
    while True: 
     yield Exception 

my_func = mock.MagicMock() 
my_func.side_effect = effect() 

my_func() #Returns "Goodbye!' 
my_func() #Raises exception 
my_func() #Raises exception 

Évidemment, vous ne voulez probablement pas d'élever un Exception nu, mais je ne sais pas ce que vous exception à vouloir soulever ...

Answer 2

Vous pouvez utiliser side_effect au lieu de return_value, par exemple:

import mock 

a = 0 
def my_func1(): 
    global a 
    a += 1 
    if a < 2: 
     return "Goodbye" 
    else: 
     raise Exception() 

my_func = mock.MagicMock() 
my_func.side_effect = my_func1 
print my_func() 
# output Goodbye 
print my_func() 
# An exception was raised.