def setUp(self):
self.ts = TargetSmartAutomation()
self.job_name = 'a-test-job'
self.sftp = SFTP(self.ts.sftp_host, os.environ['TS_SFTP_USERNAME'],
os.environ['TS_SFTP_PASSWORD'], self.ts.sftp_port)
self.test_xml = 'test/test_ts/job_config.xml'
def test_credential_validation():
with pytest.raises(ValueError):
SFTP(host=None, username=None, password=None, rsa_private_key_file=None)
with pytest.raises(ValueError):
SFTP(host=None, username='******', password='******', rsa_private_key_file='/path/to/key/file')
with pytest.raises(ValueError):
SFTP(host='host', username=None, password='******', rsa_private_key_file='/path/to/key/file')
with pytest.raises(ValueError):
SFTP(host='host', username='******', password=None, rsa_private_key_file=None)
def to_sftp_csv(self,
remote_path,
host,
username,
password,
port=22,
encoding=None,
compression=None,
errors='strict',
write_header=True,
rsa_private_key_file=None,
**csvargs):
"""
Writes the table to a CSV file on a remote SFTP server
`Args:`
remote_path: str
The remote path of the file. If it ends in '.gz', the file will be compressed.
host: str
The remote host
username: str
The username to access the SFTP server
password: str
The password to access the SFTP server
port: int
The port number of the SFTP server
encoding: str
The CSV encoding type for `csv.writer()
<https://docs.python.org/2/library/csv.html#csv.writer/>`_
errors: str
Raise an Error if encountered
write_header: boolean
Include header in output
rsa_private_key_file str
Absolute path to a private RSA key used
to authenticate stfp connection
\**csvargs: kwargs
``csv_writer`` optional arguments
""" # noqa: W605
from parsons import SFTP
sftp = SFTP(host, username, password, port, rsa_private_key_file)
compression = files.compression_type_for_path(remote_path)
local_path = self.to_csv(temp_file_compression=compression,
encoding=encoding,
errors=errors,
write_header=write_header,
**csvargs)
sftp.put_file(local_path, remote_path)
def live_sftp(simple_table, simple_csv_path, simple_compressed_csv_path):
# Generate a live SFTP connection based on these env vars
host = os.environ['SFTP_HOST']
username = os.environ['SFTP_USERNAME']
password = None
rsa_private_key_file = os.environ['SFTP_RSA_PRIVATE_KEY_FILE']
sftp = SFTP(host, username, password, rsa_private_key_file=rsa_private_key_file)
# Add a test directory and test files
sftp.make_directory(REMOTE_DIR)
sftp.put_file(simple_csv_path, REMOTE_CSV_PATH)
sftp.put_file(simple_compressed_csv_path, REMOTE_COMPRESSED_CSV_PATH)
yield sftp
# Cleanup after test
sftp.remove_file(REMOTE_CSV_PATH)
sftp.remove_file(REMOTE_COMPRESSED_CSV_PATH)
sftp.remove_directory(REMOTE_DIR)
class TestTargetSmartAutomation(unittest.TestCase):
def setUp(self):
self.ts = TargetSmartAutomation()
self.job_name = 'a-test-job'
self.sftp = SFTP(self.ts.sftp_host,
os.environ['TS_SFTP_USERNAME'],
os.environ['TS_SFTP_PASSWORD'],
self.ts.sftp_port)
self.test_xml = 'test/test_targetsmart/job_config.xml'
def tearDown(self):
# Clean up the files were put on the SFTP
self.ts.remove_files(self.job_name)
@mark_live_test
def test_create_job_xml(self):
# Assert that job xml creates the file correctly
job_xml = self.ts.create_job_xml('job_type',
'match_job',
['*****@*****.**', '*****@*****.**'])
with open(self.test_xml, 'r') as xml:
test_xml = xml.read()
with open(job_xml, 'r') as xml:
real_xml = xml.read()
self.assertEqual(test_xml, real_xml)
@mark_live_test
def test_config_status(self):
# Find good configuration
self.sftp.put_file(self.test_xml, f'{self.ts.sftp_dir}/{self.job_name}.job.xml.good')
self.assertTrue(self.ts.config_status(self.job_name))
self.ts.remove_files(self.job_name)
# Find bad configuration
self.sftp.put_file(self.test_xml, f'{self.ts.sftp_dir}/{self.job_name}.job.xml.bad')
self.assertRaises(ValueError, self.ts.config_status, self.job_name)
@mark_live_test
def test_match_status(self):
# Find good configuration
good_match = 'test/test_targetsmart/match_good.xml'
self.sftp.put_file(good_match, f'{self.ts.sftp_dir}/{self.job_name}.finish.xml')
self.assertTrue(self.ts.match_status(self.job_name))
self.ts.remove_files(self.job_name)
# Find bad configuration
bad_match = 'test/test_targetsmart/match_bad.xml'
self.sftp.put_file(bad_match, f'{self.ts.sftp_dir}/{self.job_name}.finish.xml')
self.assertRaises(ValueError, self.ts.match_status, self.job_name)
@mark_live_test
def test_remove_files(self):
# Add a file
self.sftp.put_file(self.test_xml, f'{self.ts.sftp_dir}/{self.job_name}.txt')
# Remove files
self.ts.remove_files(self.job_name)
# Check that file is not there
dir_list = self.sftp.list_directory(f'{self.ts.sftp_dir}/')
self.assertNotIn(f'{self.job_name}.txt', dir_list)
def test_credential_validation():
with pytest.raises(ValueError):
SFTP(host=None, username=None, password=None)
with pytest.raises(ValueError):
SFTP(host=None, username='******', password='******')
def generate_live_sftp_connection():
host = os.environ['SFTP_HOST']
username = os.environ['SFTP_USERNAME']
password = os.environ['SFTP_PASSWORD']
return SFTP(host, username, password)