test_mark.py

from collections import defaultdict
from dataclasses import dataclass, field

import pytest
from _pytest.mark import Mark, ParameterSet
from _pytest.python import Metafunc


@dataclass
class Metafunc:
    function = None
    config = None
    definition = None


@pytest.fixture
def metafunc():
    return Metafunc()


def parametrize_indirect(metafunc, mark):
    from _pytest.mark import ParameterSet, Mark

    if mark.name != "parametrize":
        return

    _args, _params = ParameterSet._for_parametrize(
        *mark.args, metafunc.function, metafunc.config, metafunc.definition,
    )

    @dataclass
    class Context:
        direct: list = field(default_factory=list)
        indirect: dict = field(default_factory=lambda: defaultdict(dict))

    transforms = []
    args = []
    indirects = {}

    for arg in _args:
        try:
            prefix, name = arg.split(".", 1)
            indirects[prefix] = True
            transforms.append(
                lambda ctx, val, prefix=prefix, name=name: ctx.indirect[
                    prefix
                ].setdefault(name, val)
            )
        except ValueError:
            args.append(arg)
            transforms.append(lambda ctx, val: ctx.direct.append(val))

    def generate_params():
        for param in _params:
            context = Context()

            for transform, value in zip(transforms, param.values):
                transform(context, value)

            yield ParameterSet(
                [*context.direct, *context.indirect.values()],
                marks=param.marks,
                id=param.id,
            )

    return Mark(
        mark.name,
        ((*args, *indirects.keys()), list(generate_params())),
        kwargs=dict(indirect=list(indirects.keys())),
    )


@pytest.mark.parametrize(
    ("before", "after"),
    [
        (
            Mark(
                name="parametrize",
                args=(
                    ("foo.a", "b", "foo.c", "d", "bar.e", "f"),
                    [
                        ParameterSet(
                            values=["a", "b", "c", "d", "e", "f"], marks=[], id=None
                        ),
                        ParameterSet(
                            values=["A", "B", "C", "D", "E", "F"], marks=[], id=None
                        ),
                    ],
                ),
                kwargs={},
            ),
            Mark(
                name="parametrize",
                args=(
                    ("b", "d", "f", "foo", "bar"),
                    [
                        ParameterSet(
                            values=["b", "d", "f", dict(a="a", c="c"), dict(e="e")],
                            marks=[],
                            id=None,
                        ),
                        ParameterSet(
                            values=["B", "D", "F", dict(a="A", c="C"), dict(e="E")],
                            marks=[],
                            id=None,
                        ),
                    ],
                ),
                kwargs=dict(indirect=["foo", "bar"],),
            ),
        )
    ],
)
def test_parametrize_indirect(metafunc, before, after):
    assert parametrize_indirect(metafunc, before) == after