Mock subprocess.Popen with context manager

mock_popen.py

import subprocess

from contextlib import contextmanager
from mock import Mock, MagicMock


@contextmanager
def popen_mock():
    """
    Context manager to mock subprocess calls.

    Example:

    ```
    with popen_mock() as process_mock:
       call_some_fn_that_uses_popen()
       assert process_mock.communicate.called
    ```
    """
    orig_popen = subprocess.Popen
    subprocess.Popen = MagicMock()
    process_mock = Mock()
    process_mock.configure_mock(**{
        'communicate.return_value': ('ouput', 0),
        'wait.return_value': 0,
        'returncode': 0,
    })
    subprocess.Popen.return_value = process_mock

    yield process_mock

    subprocess.Popen = orig_popen