sql-parser/run_sql_parser_test.py

from pathlib import Path
import asyncio.subprocess as subprocess
import asyncio
from watchfiles import awatch
from termcolor import colored
from datetime import datetime
import orjson
import os
import tempfile
import importlib


async def run_and_output(*args: str, timeout=10) -> tuple[bytes, bytes]:
    p = await subprocess.create_subprocess_exec(
        *args,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = await asyncio.wait_for(p.communicate(), timeout=timeout)
    return stdout, stderr


async def rebuild() -> bool:
    print(datetime.now(), colored("rebuild...", "grey"))
    stdout, _ = await run_and_output("xmake")
    if b"error" in stdout:
        print(stdout.decode("utf-8"))
        print(datetime.now(), "-" * 40)
        return False
    else:
        return True


async def assert_sql(sql: str, expected: dict):
    stdout, stderr = await run_and_output("xmake", "run", "sql-parser", sql)

    if b"error" in stdout:
        print(stdout.decode("utf-8"))
        print(datetime.now(), "-" * 40)
        print(f'other: {colored(stderr.decode("utf-8"), "yellow")}')
        assert False, "sql-parser error"

    try:
        output = orjson.loads(stdout)
    except Exception as e:
        output = {"error": e, "output": stdout.decode("utf-8")}
    open("/tmp/temp/test.py", "wb").write(
        f'"{sql}"\n\n'.encode("utf-8")
        + orjson.dumps(output, option=orjson.OPT_INDENT_2)
        + (b"\n\n" + stderr).replace(b"\n", b"\n# ")
    )
    assert (
        output == expected
    ), f"""{colored("sql-parser error", "red")}
input: {colored(sql, "yellow")}
expect: {colored(expected, "green")}
actual: {colored(output, "red")}
other: {colored(stderr.decode("utf-8"), "yellow")}

"""


async def assert_sqls():
    import test_configs.sql_parser
    importlib.reload(test_configs.sql_parser)

    for sql, excepted in test_configs.sql_parser.sql_parser_tests:
        await assert_sql(sql, excepted)


async def on_parser_modified():
    print(datetime.now(), colored("run parser tests...", "yellow"))

    try:
        await assert_sqls()
    except Exception as e:
        print(e)
    else:
        print(datetime.now(), colored("all parser tests right!", "green"))


async def assert_checks():
    import test_configs.sql_checker
    importlib.reload(test_configs.sql_checker)

    for sql, res in test_configs.sql_checker.sql_checker_tests:
        stdout, stderr = await run_and_output("xmake", "run", "sql-checker", "-s", sql)
        print(sql, res)
        if res is True:
            assert b"error" not in stdout, (
                stdout.decode("utf-8") + "\n" + stderr.decode("utf-8")
            )
            assert b"error" not in stderr, (
                stdout.decode("utf-8") + "\n" + stderr.decode("utf-8")
            )
        elif isinstance(res, str):
            res = res.encode("utf-8")
            assert res in stderr, stderr.decode("utf-8")
        else:
            assert False, f"{res} 不是合适的结果"


async def on_checker_modified():
    print(datetime.now(), colored("run checker tests...", "yellow"))
    try:
        await assert_checks()
    except Exception as e:
        print(colored(e, "red"))
    else:
        print(datetime.now(), colored("all checker tests right!", "green"))


async def assert_sql_optimizer_check():
    import test_configs.sql_optimizer
    importlib.reload(test_configs.sql_optimizer)

    for sql, res in test_configs.sql_optimizer.sql_optimizer_tests:
        stdout, stderr = await run_and_output(
            "xmake", "run", "sql-optimizer", "-s", sql
        )
        print(sql)

        try:
            output = orjson.loads(stdout)
            if output != res:
                print('', colored(res, "yellow"))
                print('real', colored(output, "red"))
                print(colored(stdout.decode("utf-8"), "yellow"))
                print(colored(stderr.decode("utf-8"), "yellow"))
                assert False
        except Exception as e:
            print(e)
            break


async def on_optimizer_modified():
    print(datetime.now(), colored("run optimizer tests...", "yellow"))
    try:
        await assert_sql_optimizer_check()
    except Exception as e:
        print(colored(e, "red"))
    else:
        print(datetime.now(), colored("all optimizer tests right!", "green"))

async def restart():
    async for _ in awatch(__file__):
        print("restart")
        os.execl("/bin/python", Path(__file__).as_posix(), Path(__file__).as_posix())


async def watch_parser():
    async for changes in awatch(
        "./src/parser.y", "./src/parser.l", "test_configs/sql_parser.py"
    ):
        if await rebuild():
            await asyncio.wait_for(on_parser_modified(), 10)


async def watch_checker():
    async for changes in awatch(
        "./src/checker.cpp",
        "./src/checker.h",
        "./src/utils.h",
        "./src/utils.cpp",
        "test_configs/sql_checker.py",
    ):
        if await rebuild():
            await on_checker_modified()


async def watch_optimizer():
    async for changes in awatch(
        "src/optimizer.cpp",
        "./src/utils.h",
        "./src/utils.cpp",
        "test_configs/sql_optimizer.py",
    ):
        if await rebuild():
            await on_optimizer_modified()


async def rerun_tests():
    await asyncio.gather(
        on_parser_modified(),
        on_checker_modified(),
    )
    await on_optimizer_modified()


async def main():
    await asyncio.gather(
        restart(),
        watch_parser(),
        watch_checker(),
        watch_optimizer(),
        rerun_tests(),
    )


if __name__ == "__main__":
    asyncio.run(main())