Mock de una intancia de un cursor de base de datos en Python

Un problema que me encontré hace poco, fue tratar de mockear una función que presenta el siguiente compartamiento. Realizar una consulta a una base de datos, el resultado es un cursor, el cual se itera con los registros que trajo. Cada iteración sobre este, se escriben los resultados en un archivo. El test que diseñado, consistía en validar si está función escribía o no este.

Un ejemplo en código de la primera aproximación del test:

output_path = tmp_path / 'output_csv.csv'

# The folling function retrive a db connection instance
conn = get_connection_db()
cur = conn.cursor.return_value
cur.fetchone.return_value = [1,2,3,4]

# Function to test
process_query_result(cur, output_path)
assert os.stat(csv_path).st_size > 0

El test busca mockear una conexión a una base de datos, así como los valores que contrendrá el cursor para luego ser iterarados y escritos en el archivo. Para esto, utilizo la propiedad return_value de mock, la cual devolverá la lista de enteros cuando cursor sea consumido desde process_query_result.

Después de ejecutar el test descripto anteriormente, falló rotundamente y tuve que interrumpir la ejecución del test porque nunca terminaba. Esto dado a como está escrita la función process_query_result (itera sando un while sobre cursor), se quedaba leyendo el cursor al infinito.

Con un poco de debug, llegué a encontrar cuál era mi problema. La propiedad return_value devolvía la lista completa de items, cada vez que se realiza una llamada la función fetchone del cursor y no un valor por vez. Entonces núnca se llegaba al fin de lista lanzando un StopIteration. En este caso, la propiedad return_value no era la solución para mi problema, así como el millón de cosas que probé buscando en muchas fuentes.

Lo que necesitaba desde un principio fue haber usado side_effect, esta propiedad mas o menos se comporta de la siguiente forma:

This can either be a function to be called when the mock is called, an iterable or an exception (class or instance) to be raised.

Justo lo que necesitaba, algo que se comporte como un iterable!, y no que me devuelva un único valor por cada vez que se ejecutaba fetchone. Entonces cambiando solo la línea siguiente:

cur.fetchone.side_effect = return_fav_list9tt

Mi test funcionó correctamente. En resúmen, si se necesita mockear un iterable que es utilizado por otra función, la mejor opción será usar side_effect.