napalm-device-types/tests/test_channel.py

"""The command channel: how a driver runs a command and opens a byte stream.

netOrk's container runtime driver speaks the Docker Engine API over a stream
from ``docker system dial-stdio`` (NetOrk/netork#765). The driver only provides
the way there; these are the SSH pieces every SSH driver can reuse. The
paramiko channel is faked, so the tests pin behaviour, not paramiko.
"""

from __future__ import annotations

import socket

import pytest

from napalm_device_types.channel import (
    CommandResult,
    ParamikoExecStream,
    open_stream_on_transport,
    run_on_transport,
)


class FakeChannel:
    """Just enough of paramiko.Channel: scripted stdout/stderr chunks, an exit code."""

    def __init__(self, stdout=(), stderr=(), exit_code=0, hang=False):
        self.out = list(stdout)
        self.err = list(stderr)
        self.exit_code = exit_code
        self.hang = hang
        self.sent = b""
        self.command = None
        self.write_closed = False
        self.closed = False
        self.timeout = None

    def exec_command(self, command):
        self.command = command

    def settimeout(self, timeout):
        self.timeout = timeout

    def sendall(self, data):
        self.sent += data

    def shutdown_write(self):
        self.write_closed = True

    def close(self):
        self.closed = True

    def recv_ready(self):
        return bool(self.out)

    def recv(self, n):
        if self.out:
            return self.out.pop(0)
        if self.hang:
            raise socket.timeout()
        return b""

    def recv_stderr_ready(self):
        return bool(self.err)

    def recv_stderr(self, n):
        return self.err.pop(0) if self.err else b""

    def exit_status_ready(self):
        return not self.hang and not self.out and not self.err

    def recv_exit_status(self):
        return self.exit_code


class FakeTransport:
    def __init__(self, channel):
        self.channel = channel

    def open_session(self):
        return self.channel


def test_run_collects_stdout_stderr_and_the_exit_code():
    ch = FakeChannel(stdout=[b"hel", b"lo\n"], stderr=[b"warn\n"], exit_code=3)

    result = run_on_transport(FakeTransport(ch), "echo hello", timeout=5)

    assert result == CommandResult(stdout="hello\n", stderr="warn\n", exit_code=3)
    assert ch.command == "echo hello"
    assert ch.closed


def test_run_sends_stdin_and_then_closes_the_write_side():
    ch = FakeChannel(stdout=[b"ok"])

    run_on_transport(FakeTransport(ch), "sudo -S true", stdin=b"pw\n", timeout=5)

    assert ch.sent == b"pw\n"
    assert ch.write_closed


def test_run_raises_when_the_command_does_not_finish_in_time():
    ch = FakeChannel(hang=True)

    with pytest.raises(TimeoutError, match="did not finish"):
        run_on_transport(FakeTransport(ch), "sleep 999", timeout=0.05)
    assert ch.closed


def test_a_stream_reads_until_eof_and_drains_stderr_on_the_way():
    ch = FakeChannel(stdout=[b"HTTP/1.1 200 OK\r\n"], stderr=[b"note\n"])
    stream = open_stream_on_transport(FakeTransport(ch), "docker system dial-stdio")

    assert stream.read(4096, timeout=5) == b"HTTP/1.1 200 OK\r\n"
    assert stream.read(4096, timeout=5) == b""
    assert stream.stderr == "note\n"
    assert ch.command == "docker system dial-stdio"


def test_a_stream_writes_and_closes_its_write_side_only_when_asked():
    """dial-stdio answers HTTP 499 to a request whose writer closed early
    (netork#771), so write() never implies close_write()."""
    ch = FakeChannel()
    stream = open_stream_on_transport(FakeTransport(ch), "docker system dial-stdio")

    stream.write(b"GET /_ping HTTP/1.1\r\n\r\n")
    assert not ch.write_closed

    stream.close_write()
    assert ch.write_closed and ch.sent.startswith(b"GET /_ping")


def test_a_stream_sends_its_stdin_prefix_first():
    """How a sudo password reaches `sudo -S` before the stream's own bytes."""
    ch = FakeChannel()

    open_stream_on_transport(FakeTransport(ch), "sudo -S -p '' cmd", stdin_prefix=b"pw\n")

    assert ch.sent == b"pw\n"


def test_a_stream_times_out_as_timeout_error():
    stream = ParamikoExecStream(FakeChannel(hang=True))

    with pytest.raises(TimeoutError):
        stream.read(10, timeout=0.01)


def test_the_stream_keeps_only_a_bounded_tail_of_stderr():
    ch = FakeChannel(stderr=[b"x" * 100, b"y" * 100])
    stream = ParamikoExecStream(ch, stderr_limit=50)

    assert stream.stderr == "y" * 50


def test_exit_status_is_none_while_running_and_the_code_after():
    running = ParamikoExecStream(FakeChannel(hang=True))
    finished = ParamikoExecStream(FakeChannel(exit_code=1))

    assert running.exit_status is None
    assert finished.exit_status == 1


def test_close_closes_the_channel():
    ch = FakeChannel()

    ParamikoExecStream(ch).close()

    assert ch.closed


def test_run_keeps_output_that_arrives_after_the_exit_status():
    """sshd may send exit-status before the last stdout bytes."""

    class EarlyExit(FakeChannel):
        def exit_status_ready(self):
            return True

        def recv_ready(self):
            return False

    channel = EarlyExit(stdout=[b"head ", b"tail"], stderr=[b"warn"])

    result = run_on_transport(FakeTransport(channel), "cmd")

    assert result == CommandResult("head tail", "warn", 0)