pytest monkeypatch mock stub asyncio.create_subprocess_exec

stubs.py

import asyncio
import re
from dataclasses import dataclass
from os import PathLike
from typing import List


class StubAsyncProcess(asyncio.subprocess.Process):
    predefined_returncode = 0
    predefined_communicate_return = (b"", b"")

    def __init__(self, *args, **kwargs):
        pass

    @property
    def returncode(self):
        return StubAsyncProcess.predefined_returncode

    async def communicate(self, *args, **kwargs):
        return StubAsyncProcess.predefined_communicate_return


@dataclass
class StubCommand:
    pattern: str = ".*"
    return_code: int = 0
    stdout: bytes = b""
    stderr: bytes = b""


class StubAsyncExec:
    def __init__(self, commands: List[StubCommand]):
        """Factory for generating stub for asyncio.create_subprocess_exec based on command patterns

        Args:
            commands: List of commands to be replaced with a stub
        """
        
        self.commands = commands
        
        # Store original create_subprocess_exec before it gets monkey patched
        self.original_create_subprocess_exec = asyncio.create_subprocess_exec        

    def create_subprocess_exec(self):
        async def stub(program: str | bytes | PathLike, *args: str | bytes | PathLike, **kwargs):
            command_string = "".join([program, *args])
            for command in self.commands:
                regex = re.compile(command.pattern)
                
                # Intercept command that matches pattern and defines the expected return code, stdout and stderr
                if regex.match(command_string):
                    StubAsyncProcess.predefined_returncode = command.return_code
                    StubAsyncProcess.predefined_communicate_return = (command.stdout, command.stderr)
                    return StubAsyncProcess()

            # If no commands are matched, use original asyncio.create_subprocess_exec
            return await self.original_create_subprocess_exec(program, *args, **kwargs)

        return stub

test_async_exec.py

import asyncio
from dataclasses import dataclass

import pytest

from .stubs import StubAsyncExec, StubCommand


@dataclass
class Result:
    return_code: int
    stdout: bytes = None
    stderr: bytes = None


async def exec_cmd(*cmd: str) -> Result:
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await proc.communicate()
    return Result(proc.returncode, stdout, stderr)


@pytest.mark.asyncio
async def test_async_exec_stub(monkeypatch):
    stub = StubAsyncExec(
        [
            StubCommand(pattern="date.*", stdout=b"It is today!", return_code=0),
            StubCommand(pattern="ls.*", stderr=b"no no no", return_code=1),
        ]
    )
    monkeypatch.setattr("asyncio.create_subprocess_exec", stub.create_subprocess_exec())

    dt = await exec_cmd("date")
    assert dt.return_code == 0
    assert dt.stdout == b"It is today!"
    assert not dt.stderr

    ls = await exec_cmd("ls", "-l")
    assert ls.return_code == 1
    assert not ls.stdout
    assert ls.stderr == b"no no no"

    # Uses original asyncio.create_subprocess_exec
    echo = await exec_cmd("echo", "1")
    assert echo.return_code == 0
    assert echo.stdout == b"1\n"
    assert not echo.stderr