PropertyMock с изменением значений после каждого вызова

Я пытаюсь издеваться над psycopg2, и мне нужно обработать 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]

Я могу сделать это через PropertyMock:

  type(fake_cursor).description = PropertyMock(return_value=descriptions[0])

но это не решает мою проблему, потому что тестируемая система выполняет два запроса, и описание будет другим во второй раз.

Как я могу изменить возвращаемое значение макета при втором вызове?

Я попытался:

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)

Но я получаю

E     UnboundLocalError: local variable 'cursor_execute_call_count' referenced before assignment

(Для меня это не имеет смысла. Должна быть странная проблема с определением объема.)

В конце концов я заставил его работать с этим кодом, но похоже, что должен быть способ лучше:

def handle_cursor_execute(arg1, arg2):
  description = descriptions.pop()
  type(fake_cursor).description = PropertyMock(return_value=description)

1 ответ

Решение

Изменить на основе уточнения:

Думаю, вы можете реализовать итератор:

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())

Вот так это будет выглядеть:

>>> 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]
>>>
Другие вопросы по тегам