Mockig a database cursor instance in Python

I found an issue recently, when I tried to mock a function that has the following behavior. There is a query to database to fetch several rows, the result is a cursor that will be iterated later. Hence each iteration will write a file with the result. The designed test assert if the function write or not the file.

A code example about my first attemp:

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

The test aims to mock both a new database connection and the cursor. Later that the cursor will be iterated and after call fetchone the result will be written as a line into the file. In this case I used the property return_value from mock. This property will return a list of integer each time that cursor is consumed from process_query_result.

The test step described before did not work also I needed to interrupt the execution due to the test never end. That situation happent because the function process_query_result use a while statement over the cursor and never break the loop.

With a little bit of debuging time, I found which was my mistake. The property return_value returned the full list of integers each time that fetch_one function of the cursor was called. So the loop never reaches the end in consequence the iterator never got an StopIteration.

What I was looking for from the very first time was the side_effect property. This one has more or less the following behave:

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.

Eureka! what I needed! Something that mimics an iterable! From now fetchone will not reutrn the full list. Finally I needed just tweak a single line as follow:

cur.fetchone.side_effect = return_fav_list

So the test ran right this time. In a nutshell if you need mock an iterable instance inside a fuction, the best option will be use side_effect.