I’ve written a function that establishes and returns an SFTP client connection. Here is the code:

import paramiko
from io import StringIO
from fastapi import HTTPException

def get_sftp_client(
    username: str,
    password: str = None,
    key_filename: str = None,
    private_key: str = None,
    host: str = "hostexample.com",
):
    try:
        if not username or (not password and not key_filename and not private_key):
            raise HTTPException(status_code=500, detail="SFTP credentials are not set")

        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        if private_key:
            private_key_file = StringIO(private_key)
            pkey = paramiko.RSAKey.from_private_key(private_key_file)
            ssh.connect(
                hostname=host,
                port=22,
                username=username,
                pkey=pkey,
            )
        elif key_filename:
            ssh.connect(
                hostname=host,
                port=22,
                username=username,
                key_filename=key_filename,
            )
        else:
            ssh.connect(
                hostname=host,
                port=22,
                username=username,
                password=password,
            )

        sftp = ssh.open_sftp()
        return sftp
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to connect to SFTP server: {str(e)}")

I would like to write unit tests for this function, especially to understand how it behaves when both a password and a private key are provided. I’m struggling to mock the dependencies properly and I’m not sure how to test this without a local SFTP server.

Here is what I have tried so far using pytest and unittest.mock:

import pytest
from unittest.mock import patch, MagicMock
from io import StringIO

def test_get_sftp_client_password_and_private_key():
    from my_module import get_sftp_client, HTTPException 

    # Mocking paramiko
    with patch('my_module.paramiko.SSHClient') as MockSSHClient:
        mock_ssh_client = MockSSHClient.return_value
        mock_ssh_client.open_sftp.return_value = MagicMock()

        # Test with password and private_key
        username = "testuser"
        password = "testpassword"
        private_key = "test_private_key"

        # Simulate private key behavior
        mock_pkey = MagicMock()
        with patch('my_module.paramiko.RSAKey.from_private_key', return_value=mock_pkey):
            sftp_client = get_sftp_client(
                username=username,
                password=password,
                private_key=private_key
            )

        # Assertions
        MockSSHClient.assert_called_once()
        mock_ssh_client.set_missing_host_key_policy.assert_called_once()
        mock_ssh_client.connect.assert_called_once_with(
            hostname="hostexample.com",
            port=22,
            username=username,
            pkey=mock_pkey,
        )
        mock_ssh_client.open_sftp.assert_called_once()
        assert sftp_client is not None

def test_get_sftp_client_raises_exception():
    from my_module import get_sftp_client, HTTPException 

    # Mocking paramiko
    with patch('my_module.paramiko.SSHClient') as MockSSHClient:
        mock_ssh_client = MockSSHClient.return_value
        mock_ssh_client.connect.side_effect = Exception("Connection failed")

        # Test with missing credentials
        username = "testuser"

        with pytest.raises(HTTPException) as excinfo:
            get_sftp_client(username=username)

        assert excinfo.value.status_code == 500
        assert "Failed to connect to SFTP server" in excinfo.value.detail

I’m looking for advice on how to properly mock the paramiko library to test this function effectively. Any tips on how to test this without a local SFTP server would be greatly appreciated as well.

Thank you in advance for your help!