2021年6月29日火曜日

paramiko を monkeypatch する方法

概要
環境
実行コード
テストコード
解説
scp の場合
scp のテストコード
解説
おまけ: ProxyCommand を使う場合は
おまけ2: 鍵認証の場合は

概要

Python の SSH クライアントライブラリ paramiko を pytest の monkeypatch でモックする方法を紹介します with - as 句を使うので __enter__ と __exit__ を実装する必要があります

環境

macOS 11.4
Python 3.8.3
paramiko 2.7.2

実行コード

import paramiko

def execute_cmd():
    with paramiko.SSHClient() as ssh:
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname='192.168.100.99',
                    port=22,
                    username='user1',
                    password='xxxxxxxxxxx')
        stdin, stdout, stderr = ssh.exec_command('hostname')
        stdin.close()
        return stdout.read()

テストコード

import paramiko

from test import execute_cmd

class DummySSHClient():
    def __init__(self):
        self

    def set_missing_host_key_policy(self, policy):
        self.policy = policy

    def connect(self, hostname=None, port=None, username=None, password=None):
        return None

    def exec_command(self, cmd):
        return DummyStdin(), DummyStdout(), ""

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

    def close(self):
        return None

class DummyStdout():
    def read(self):
        return "testhost1"

class DummyStdin():
    def close(self):
        return None

class DummyAutoAddPolicy():
    def __init__(self):
        pass

def test_execute_cmd(monkeypatch):
    monkeypatch.setattr(paramiko, "SSHClient", DummySSHClient)
    monkeypatch.setattr(paramiko, "AutoAddPolicy", DummyAutoAddPolicy)
    result = execute_cmd()
    assert (result == "testhost1")

解説

paramiko.SSHClient をダミーの DummySSHClient でモックします paramiko.AutoAddPolicy もモックが必要なので DummyAutoAddPolicy を作成します

stdout と stdin のオブジェクトも使う場合にはそれぞれモックします使用しない場合はモック不要です

あとは monkeypatch で setatter で SSHClient と AutoAddPolicy を持っくすれば OK です

scp の場合

paramiko + scp でファイルをアップロードする場合も紹介します

import paramiko
import scp

def upload_file():
    with paramiko.SSHClient() as ssh:
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname='192.168.100.99',
                    port=22,
                    username='user1',
                    password='xxxxxxxxxxx')
        with scp.SCPClient(ssh.get_transport()) as scp:
            scp.put('hoge.txt', '/tmp/hoge.txt')

scp のテストコード

import paramiko
import scp

from test import execute_cmd, upload_file

class DummySSHClient():
    def __init__(self):
        self

    def set_missing_host_key_policy(self, policy):
        self.policy = policy

    def get_transport(self):
        return None

    def connect(self, hostname=None, port=None, username=None, password=None):
        return None

    def exec_command(self, cmd):
        return DummyStdin(), DummyStdout(), ""

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

class DummySCPClient():
    def __init__(self, transport):
        self

    def put(self, src, dict):
        return None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

class DummyStdout():
    def read(self):
        return "testhost1"

class DummyStdin():
    def close(self):
        return None

class DummyAutoAddPolicy():
    def __init__(self):
        pass

def test_upload_file(monkeypatch):
    monkeypatch.setattr(paramiko, "SSHClient", DummySSHClient)
    monkeypatch.setattr(paramiko, "AutoAddPolicy", DummyAutoAddPolicy)
    monkeypatch.setattr(scp, "SCPClient", DummySCPClient)
    result = upload_file()
    assert (result == None)

解説

scp の場合もほぼ同じです DummySSHClient 側に get_transport を追加で実装してあげます

おまけ: ProxyCommand を使う場合は

import paramiko

def execute_cmd_with_proxy():
    with paramiko.SSHClient() as ssh:
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname='192.168.100.99',
                    port=22,
                    username='user1',
                    password='xxxxxxxxxxx',
                    sock=paramiko.ProxyCommand('nc -X connect -x 192.168.100.100:3128 192.168.100.99 22'))
        stdin, stdout, stderr = ssh.exec_command('hostname')
        stdin.close()
        print(stdout.read())

になるので connect に 1 つ変数を追加してあげれば OK です (以下途中省略)

class DummySSHClient():
    def connect(self, hostname=None, port=None, username=None, password=None, sock=None):
        return None

おまけ2: 鍵認証の場合は

class DummyRSAKey():
    @classmethod
    def from_private_key_file(cls, key_path, passphrase):
        return "key"

class DummySSHClient():
    def connect(self, hostname=None, port=None, username=None, password=None, sock=None, pkey=None):
        return None

monkeypatch.setattr(paramiko, "RSAKey", DummyRSAKey)

こんな感じです

hawksnowlog