Home > Enterprise >  Can multiple object attributes be patched at the same time?
Can multiple object attributes be patched at the same time?

Time:01-22

I am looking to reduce nesting in my tests, I was looking at with patch.multiple() for the likes of below, but can't figure out a way to get it to work.

Where mock_task_instance, mock_dag_run and success_mock_task_instance are defined in the test function.

So I was wondering if there is a way that multiple object attributes be patched at the same time?

with patch.object(
            mock_task_instance, "xcom_pull", side_effect=[self.file_name, self.config]
        ):
            with patch.object(
                mock_dag_run, "get_task_instances", return_value=[success_mock_task_instance]
            ):
                with patch.object(
                    success_mock_task_instance, "current_state", return_value=State.SUCCESS
                ):

                    _my_func()

CodePudding user response:

If you are looking for a way to cleanup nested with blocks, I recommend checking out contextlib.ExitStack() as this can help with than.

Here is the example from the docs:

with ExitStack() as stack:
    files = [stack.enter_context(open(fname)) for fname in filenames]
    # All opened files will automatically be closed at the end of
    # the with statement, even if attempts to open files later
    # in the list raise an exception

See: https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack

For your code, I think that might look a bit like:

with contextlib.ExitStack() as stack:
    stack.enter_context(patch.object(mock_task_instance, "xcom_pull", side_effect=[self.file_name, self.config]))
    stack.enter_context(patch.object(mock_dag_run, "get_task_instances", return_value=[success_mock_task_instance]))
    stack.enter_context(patch.object(success_mock_task_instance, "current_state", return_value=State.SUCCESS))
    _my_func()
  •  Tags:  
  • Related