Reputation: 4743
I'm trying to mock psycopg2 and need to handle cursor.description
# run query 1
column_names = [desc[0] for desc in cursor.description]
# run query 2
# NB cursor.description now returns different value
column_names = [desc[0] for desc in cursor.description]
I can do this via a PropertyMock:
type(fake_cursor).description = PropertyMock(return_value=descriptions[0])
but it doesn't solve my problem because the system under test does two queries and the description will be different for the second time.
How can I change the return value of the mock on the second call?
I tried:
cursor_execute_call_count = 0
def handle_cursor_execute(arg1, arg2):
cursor_execute_call_count = cursor_execute_call_count + 1
type(fake_cursor).description = PropertyMock(return_value=descriptions[cursor_execute_call_count])
fake_cursor.execute = Mock(side_effect=handle_cursor_execute)
But I get
E UnboundLocalError: local variable 'cursor_execute_call_count' referenced before assignment
(That doesn't make sense to me. There must be a weird scoping issue going on.)
I did eventually get it working with this code but it seems like there must be a better way:
def handle_cursor_execute(arg1, arg2):
description = descriptions.pop()
type(fake_cursor).description = PropertyMock(return_value=description)
Upvotes: 0
Views: 273
Reputation: 388
Edit based on clarification:
I Guess you can implement an iterator:
class MyDescription():
descriptions = iter([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
def __iter__(self):
for description_list in self.descriptions:
for description in description_list:
yield description
raise StopIteration
type(fake_cursor).description = PropertyMock(return_value=MyDescription())
So it will look like this:
>>> z = MyDescription()
>>> [i for i in z]
[1, 2, 3]
>>> [i for i in z]
[4, 5, 6]
>>> [i for i in z]
[7, 8, 9]
>>>
Upvotes: 2