from unittest import TestCase
import doctest
import pkgutil
import re

import av


def fix_doctests(suite):

    for case in suite._tests:

        # Add some more flags.
        case._dt_optionflags = (
            (case._dt_optionflags or 0)
            | doctest.IGNORE_EXCEPTION_DETAIL
            | doctest.ELLIPSIS
            | doctest.NORMALIZE_WHITESPACE
        )

        case._dt_test.globs["av"] = av
        case._dt_test.globs["video_path"] = av.datasets.curated(
            "pexels/time-lapse-video-of-night-sky-857195.mp4"
        )

        for example in case._dt_test.examples:

            # Remove b prefix from strings.
            if example.want.startswith("b'"):
                example.want = example.want[1:]


def register_doctests(mod):

    if isinstance(mod, str):
        mod = __import__(mod, fromlist=[""])

    try:
        suite = doctest.DocTestSuite(mod)
    except ValueError:
        return

    fix_doctests(suite)

    cls_name = "Test" + "".join(x.title() for x in mod.__name__.split("."))
    cls = type(cls_name, (TestCase,), {})

    for test in suite._tests:

        def func(self):
            return test.runTest()

        name = str("test_" + re.sub("[^a-zA-Z0-9]+", "_", test.id()).strip("_"))
        func.__name__ = name
        setattr(cls, name, func)

    globals()[cls_name] = cls


for importer, mod_name, ispkg in pkgutil.walk_packages(
    path=av.__path__, prefix=av.__name__ + ".", onerror=lambda x: None
):
    register_doctests(mod_name)
