test_shell.py

#!/usr/bin/env python3

import logging
import sys
import unittest
from dataclasses import dataclass
from typing import Any, cast, List

import expecttest

import ghstack.shell


@dataclass
class ConsoleMsg:
    pass


@dataclass
class out(ConsoleMsg):
    msg: str


@dataclass
class err(ConsoleMsg):
    msg: str


@dataclass
class big_dump(ConsoleMsg):
    pass


class TestShell(expecttest.TestCase, unittest.IsolatedAsyncioTestCase):
    def setUp(self) -> None:
        self.sh = ghstack.shell.Shell()
        asyncio_logger = logging.getLogger("asyncio")
        previous_asyncio_level = asyncio_logger.level
        self.addCleanup(asyncio_logger.setLevel, previous_asyncio_level)
        asyncio_logger.setLevel(logging.ERROR)

    async def emit(
        self, *payload: ConsoleMsg, **kwargs: Any
    ) -> ghstack.shell._SHELL_RET:
        args: List[str] = [sys.executable, "emitter.py"]
        for p in payload:
            if isinstance(p, out):
                args.extend(("o", p.msg))
            elif isinstance(p, err):
                args.extend(("e", p.msg))
            elif isinstance(p, big_dump):
                args.extend(("r", "-"))
        return cast(ghstack.shell._SHELL_RET, await self.sh.ash(*args, **kwargs))

    def flog(self, cm: "unittest._AssertLogsContext") -> str:  # type: ignore[name-defined]
        def redact(s: str) -> str:
            s = s.replace(sys.executable, "python")
            s = s.replace("'python'", "python")
            return s

        return "\n".join(redact(r.getMessage()) for r in cm.records)

    async def test_stdout(self) -> None:
        with self.assertLogs(level=logging.DEBUG) as cm:
            await self.emit(out(r"arf\n"))
        self.assertExpectedInline(
            self.flog(cm),
            """\
$ python emitter.py o 'arf\\n'
arf
""",
        )

    async def test_stderr(self) -> None:
        with self.assertLogs(level=logging.DEBUG) as cm:
            await self.emit(err(r"arf\n"))
        self.assertExpectedInline(
            self.flog(cm),
            """\
$ python emitter.py e 'arf\\n'
# stderr:
arf
""",
        )

    async def test_stdout_passthru(self) -> None:
        with self.assertLogs(level=logging.DEBUG) as cm:
            await self.emit(out(r"arf\n"), stdout=None)
        self.assertExpectedInline(
            self.flog(cm),
            """\
$ python emitter.py o 'arf\\n'
arf
""",
        )

    async def test_stdout_with_stderr_prefix(self) -> None:
        # What most commands should look like
        with self.assertLogs(level=logging.DEBUG) as cm:
            await self.emit(
                err(r"Step 1...\n"),
                err(r"Step 2...\n"),
                err(r"Step 3...\n"),
                out(r"out\n"),
                stdout=None,
            )
        self.assertExpectedInline(
            self.flog(cm),
            """\
$ python emitter.py e 'Step 1...\\n' e 'Step 2...\\n' e 'Step 3...\\n' o 'out\\n'
# stderr:
Step 1...
Step 2...
Step 3...

# stdout:
out
""",
        )

    async def test_interleaved_stdout_stderr_passthru(self) -> None:
        # NB: stdout is flushed in each of these cases
        with self.assertLogs(level=logging.DEBUG) as cm:
            await self.emit(
                out(r"A\n"), err(r"B\n"), out(r"C\n"), err(r"D\n"), stdout=None
            )
        self.assertExpectedInline(
            self.flog(cm),
            """\
$ python emitter.py o 'A\\n' e 'B\\n' o 'C\\n' e 'D\\n'
# stderr:
B
D

# stdout:
A
C
""",
        )

    async def test_deadlock(self) -> None:
        await self.emit(big_dump())

    async def test_uses_raw_fd(self) -> None:
        await self.emit(out(r"A\n"), stdout=sys.stdout)


if __name__ == "__main__":
    unittest.main()