asciinema/tests/pty_test.py

import os
import pty

import asciinema.pty_

from .test_helper import Test


class FakeStdout:
    def __init__(self):
        self.data = []

    def write_stdout(self, ts, data):
        self.data.append(data)

    def write_stdin(self, ts, data):
        pass


class TestRecord(Test):
    def setUp(self):
        self.real_os_write = os.write
        os.write = self.os_write

    def tearDown(self):
        os.write = self.real_os_write

    def os_write(self, fd, data):
        if fd != pty.STDOUT_FILENO:
            self.real_os_write(fd, data)

    def test_record_command_writes_to_stdout(self):
        output = FakeStdout()

        command = [
            "python3",
            "-c",
            (
                "import sys"
                "; import time"
                "; sys.stdout.write('foo')"
                "; sys.stdout.flush()"
                "; time.sleep(0.01)"
                "; sys.stdout.write('bar')"
            ),
        ]
        asciinema.pty_.record(command, output, lambda: (80, 24))

        assert output.data == [b"foo", b"bar"]