<s> from unittest import TestCase <EOL> from socket import gaierror , error as socket_error <EOL> from mock import patch , mock_open <EOL> from paramiko import SSHException , AuthenticationException <EOL> from pyinfra . api import ssh <EOL> from . paramiko_util import ( <EOL> FakeSSHClient , FakeSFTPClient , FakeRSAKey , <EOL> FakeAgentRequestHandler <EOL> ) <EOL> ssh . SSHClient = FakeSSHClient <EOL> ssh . SFTPClient = FakeSFTPClient <EOL> ssh . RSAKey = FakeRSAKey <EOL> ssh . AgentRequestHandler = FakeAgentRequestHandler <EOL> from pyinfra . api import Inventory , Config , State <EOL> from pyinfra . api . ssh import connect_all , connect <EOL> from pyinfra . api . operation import add_op <EOL> from pyinfra . api . operations import run_ops <EOL> from pyinfra . api . exceptions import PyinfraError <EOL> from pyinfra . modules import files , server <EOL> from . util import create_host <EOL> def make_inventory ( hosts = ( '<STR_LIT>' , '<STR_LIT>' ) , ** kwargs ) : <EOL> return Inventory ( <EOL> ( hosts , { } ) , <EOL> test_group = ( [ <EOL> '<STR_LIT>' <EOL> ] , { <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } ) , <EOL> ssh_user = '<STR_LIT>' , <EOL> ssh_key = '<STR_LIT:test>' , <EOL> ** kwargs <EOL> ) <EOL> def make_config ( FAIL_PERCENT = <NUM_LIT:0> , TIMEOUT = <NUM_LIT:1> , ** kwargs ) : <EOL> return Config ( <EOL> FAIL_PERCENT = FAIL_PERCENT , <EOL> TIMEOUT = TIMEOUT , <EOL> ** kwargs <EOL> ) <EOL> class TestApi ( TestCase ) : <EOL> def test_inventory_creation ( self ) : <EOL> inventory = make_inventory ( ) <EOL> host = inventory [ '<STR_LIT>' ] <EOL> self . assertEqual ( host . data . ssh_user , '<STR_LIT>' ) <EOL> self . assertEqual ( <EOL> inventory . get_groups_data ( [ '<STR_LIT>' ] ) . dict ( ) , <EOL> { <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } <EOL> ) <EOL> def test_connect_all ( self ) : <EOL> inventory = make_inventory ( ) <EOL> state = State ( inventory , make_config ( ) ) <EOL> connect_all ( state ) <EOL> self . assertEqual ( len ( inventory . connected_hosts ) , <NUM_LIT:2> ) <EOL> def test_connect_all_password ( self ) : <EOL> inventory = make_inventory ( ssh_password = '<STR_LIT:test>' ) <EOL> state = State ( inventory , make_config ( ) ) <EOL> connect_all ( state ) <EOL> self . assertEqual ( len ( inventory . connected_hosts ) , <NUM_LIT:2> ) <EOL> def test_connect_exceptions_fail ( self ) : <EOL> for exception in ( <EOL> AuthenticationException , SSHException , <EOL> gaierror , socket_error <EOL> ) : <EOL> host = create_host ( name = '<STR_LIT>' , data = { <EOL> '<STR_LIT>' : AuthenticationException <EOL> } ) <EOL> self . assertEqual ( connect ( host ) , None ) <EOL> def test_fail_percent ( self ) : <EOL> inventory = make_inventory ( ( '<STR_LIT>' , SSHException ) ) <EOL> state = State ( inventory , make_config ( ) ) <EOL> with self . assertRaises ( PyinfraError ) : <EOL> connect_all ( state ) <EOL> def test_basic_op ( self ) : <EOL> state = State ( make_inventory ( ) , make_config ( ) ) <EOL> connect_all ( state ) <EOL> add_op ( <EOL> state , files . file , <EOL> '<STR_LIT>' , <EOL> user = '<STR_LIT>' , <EOL> group = '<STR_LIT>' , <EOL> mode = '<STR_LIT>' , <EOL> sudo = True <EOL> ) <EOL> run_ops ( state ) <EOL> def test_file_op ( self ) : <EOL> state = State ( make_inventory ( ) , make_config ( ) ) <EOL> connect_all ( state ) <EOL> with patch ( '<STR_LIT>' , mock_open ( read_data = '<STR_LIT>' ) , create = True ) : <EOL> add_op ( <EOL> state , files . put , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' <EOL> ) <EOL> with patch ( '<STR_LIT>' , mock_open ( read_data = '<STR_LIT>' ) , create = True ) : <EOL> add_op ( <EOL> state , files . put , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> sudo = True , <EOL> sudo_user = '<STR_LIT>' <EOL> ) <EOL> run_ops ( state ) <EOL> def test_run_ops ( self ) : <EOL> state = State ( make_inventory ( ) , make_config ( ) ) <EOL> connect_all ( state ) <EOL> add_op ( <EOL> state , server . shell , <EOL> '<STR_LIT>' <EOL> ) <EOL> run_ops ( state ) <EOL> def test_run_ops_serial ( self ) : <EOL> state = State ( make_inventory ( ) , make_config ( ) ) <EOL> connect_all ( state ) <EOL> add_op ( <EOL> state , server . shell , <EOL> '<STR_LIT>' <EOL> ) <EOL> run_ops ( state , serial = True ) <EOL> def test_run_ops_no_wait ( self ) : <EOL> state = State ( make_inventory ( ) , make_config ( ) ) <EOL> connect_all ( state ) <EOL> add_op ( <EOL> state , server . shell , <EOL> '<STR_LIT>' <EOL> ) <EOL> run_ops ( state , no_wait = True ) </s>
