def test_wait_for_transfer_job(self, mock_list, mock_sleep):
mock_list.side_effect = [
[{
METADATA: {
STATUS: GcpTransferOperationStatus.IN_PROGRESS
}
}],
[{
METADATA: {
STATUS: GcpTransferOperationStatus.SUCCESS
}
}],
]
job_name = 'transferJobs/test-job'
self.gct_hook.wait_for_transfer_job({
PROJECT_ID: TEST_PROJECT_ID,
'name': job_name
})
calls = [
mock.call(request_filter={
FILTER_PROJECT_ID: TEST_PROJECT_ID,
FILTER_JOB_NAMES: [job_name]
}),
mock.call(request_filter={
FILTER_PROJECT_ID: TEST_PROJECT_ID,
FILTER_JOB_NAMES: [job_name]
})
]
mock_list.assert_has_calls(calls, any_order=True)
mock_sleep.assert_called_once_with(TIME_TO_SLEEP_IN_SECONDS)
def test_execute_move_with_wildcard(self, sftp_hook, gcs_hook):
sftp_hook.return_value.get_tree_map.return_value = [
[
"main_dir/test_object3.json",
"main_dir/sub_dir/test_object3.json"
],
[],
[],
]
gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2]
task = SFTPToGoogleCloudStorageOperator(
task_id=TASK_ID,
source_path=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=TEST_BUCKET,
destination_path=DESTINATION_PATH_DIR,
move_object=True,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
task.execute(None)
sftp_hook.return_value.delete_file.assert_has_calls([
mock.call("main_dir/test_object3.json"),
mock.call("main_dir/sub_dir/test_object3.json"),
])
def test_sensor(self, mock_describe, hook_init, mock_get_conn):
hook_init.return_value = None
mock_describe.side_effect = [
DESCRIBE_ENDPOINT_CREATING_RESPONSE,
DESCRIBE_ENDPOINT_UPDATING_RESPONSE,
DESCRIBE_ENDPOINT_INSERVICE_RESPONSE
]
sensor = SageMakerEndpointSensor(task_id='test_task',
poke_interval=1,
aws_conn_id='aws_test',
endpoint_name='test_job_name')
sensor.execute(None)
# make sure we called 3 times(terminated when its completed)
self.assertEqual(mock_describe.call_count, 3)
# make sure the hook was initialized with the specific params
calls = [
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test')
]
hook_init.assert_has_calls(calls)
def test_execute(self, gcs_mock_hook, s3_one_mock_hook, s3_two_mock_hook):
"""Test the execute function when the run is successful."""
operator = S3ToGoogleCloudStorageOperator(task_id=TASK_ID,
bucket=S3_BUCKET,
prefix=S3_PREFIX,
delimiter=S3_DELIMITER,
dest_gcs_conn_id=GCS_CONN_ID,
dest_gcs=GCS_PATH_PREFIX)
s3_one_mock_hook.return_value.list_keys.return_value = MOCK_FILES
s3_two_mock_hook.return_value.list_keys.return_value = MOCK_FILES
uploaded_files = operator.execute(None)
gcs_mock_hook.return_value.upload.assert_has_calls([
mock.call('gcs-bucket', 'data/TEST1.csv', mock.ANY, gzip=False),
mock.call('gcs-bucket', 'data/TEST3.csv', mock.ANY, gzip=False),
mock.call('gcs-bucket', 'data/TEST2.csv', mock.ANY, gzip=False)
],
any_order=True)
s3_one_mock_hook.assert_called_once_with(aws_conn_id=AWS_CONN_ID,
verify=None)
s3_two_mock_hook.assert_called_once_with(aws_conn_id=AWS_CONN_ID,
verify=None)
gcs_mock_hook.assert_called_once_with(
google_cloud_storage_conn_id=GCS_CONN_ID, delegate_to=None)
# we expect MOCK_FILES to be uploaded
self.assertEqual(sorted(MOCK_FILES), sorted(uploaded_files))
def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock_sync):
executor = BaseExecutor()
executor.heartbeat()
calls = [mock.call('executor.open_slots', mock.ANY),
mock.call('executor.queued_tasks', mock.ANY),
mock.call('executor.running_tasks', mock.ANY)]
mock_stats_gauge.assert_has_calls(calls)
def test_write_local_setting(self):
"""Establish that if we attempt to write a valid setting locally, that
the correct parser's write method is run.
"""
# Invoke the command, but trap the file-write at the end
# so we don't plow over real things.
mock_open = mock.mock_open()
with mock.patch('tower_cli.cli.misc.open', mock_open,
create=True):
with mock.patch.object(os, 'chmod') as chmod:
result = self.runner.invoke(
config, ['username', 'meagan', '--scope=local'],
)
filename = ".tower_cli.cfg"
chmod.assert_called_once_with(filename, int('0600', 8))
# Ensure that the command completed successfully.
self.assertEqual(result.exit_code, 0)
self.assertEqual(result.output.strip(),
'Configuration updated successfully.')
# Ensure that the output seems to be correct.
self.assertIn(mock.call('.tower_cli.cfg', 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('username = meagan\n'),
mock_open.mock_calls)
def test_write_global_setting(self):
"""Establish that if we attempt to write a valid setting, that
the parser's write method is run.
"""
# Invoke the command, but trap the file-write at the end
# so we don't plow over real things.
filename = '/etc/tower/tower_cli.cfg'
mock_open = mock.mock_open()
with mock.patch('tower_cli.commands.config.open',
mock_open,
create=True):
with mock.patch.object(os.path, 'isdir') as isdir:
with mock.patch.object(os, 'chmod') as chmod:
isdir.return_value = True
result = self.runner.invoke(
config,
['username', 'luke', '--scope=global'],
)
isdir.assert_called_once_with('/etc/tower/')
chmod.assert_called_once_with(filename, int('0600', 8))
# Ensure that the command completed successfully.
self.assertEqual(result.exit_code, 0)
self.assertEqual(result.output.strip(),
'Configuration updated successfully.')
# Ensure that the output seems to be correct.
self.assertIn(mock.call('/etc/tower/tower_cli.cfg', 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('username = luke\n'),
mock_open.mock_calls)
def test_response_in_logs_after_failed_check(self, m):
"""
Test that when using SimpleHttpOperator with log_response=True,
the response is logged even if request_check fails
"""
def response_check(response):
return response.text != 'invalid response'
m.get('http://www.example.com', text='invalid response')
operator = SimpleHttpOperator(
task_id='test_HTTP_op',
method='GET',
endpoint='/',
http_conn_id='HTTP_EXAMPLE',
log_response=True,
response_check=response_check
)
with mock.patch.object(operator.log, 'info') as mock_info:
self.assertRaises(AirflowException, operator.execute, None)
calls = [
mock.call('Calling HTTP method'),
mock.call('invalid response')
]
mock_info.assert_has_calls(calls, any_order=True)
def test_should_copy_single_file(self, mock_named_temporary_file,
mock_gdrive, mock_gcs_hook):
type(mock_named_temporary_file.return_value.__enter__.return_value
).name = mock.PropertyMock(side_effect=["TMP1"])
task = GcsToGDriveOperator(
task_id="copy_single_file",
source_bucket="data",
source_object="sales/sales-2017/january.avro",
destination_object="copied_sales/2017/january-backup.avro",
)
task.execute(mock.MagicMock())
mock_gcs_hook.assert_has_calls([
mock.call(delegate_to=None,
google_cloud_storage_conn_id="google_cloud_default"),
mock.call().download(bucket="data",
filename="TMP1",
object="sales/sales-2017/january.avro"),
])
mock_gdrive.assert_has_calls([
mock.call(delegate_to=None, gcp_conn_id="google_cloud_default"),
mock.call().upload_file(
local_location="TMP1",
remote_location="copied_sales/2017/january-backup.avro"),
])
def test_write_global_setting(self):
"""Establish that if we attempt to write a valid setting, that
the parser's write method is run.
"""
# Invoke the command, but trap the file-write at the end
# so we don't plow over real things.
mock_open = mock.mock_open()
with mock.patch('tower_cli.commands.config.open', mock_open,
create=True):
with mock.patch.object(os.path, 'isdir') as isdir:
isdir.return_value = True
result = self.runner.invoke(config,
['username', 'luke', '--scope=global'],
)
isdir.assert_called_once_with('/etc/awx/')
# Ensure that the command completed successfully.
self.assertEqual(result.exit_code, 0)
self.assertEqual(result.output.strip(),
'Configuration updated successfully.')
# Ensure that the output seems to be correct.
self.assertIn(mock.call('/etc/awx/tower_cli.cfg', 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('username = luke\n'),
mock_open.mock_calls)
def test_personal_access_token(self):
"""Establish that if `tower-cli login` is called with a username and
password, we obtain and write an oauth token to the config file
"""
# Invoke the command.
mock_open = mock.mock_open()
with mock.patch('tower_cli.cli.misc.open', mock_open,
create=True):
with mock.patch.object(os, 'chmod'):
with client.test_mode as t:
# You have to modify this internal private registry to
# register a URL endpoint that _doesn't_ have the version
# prefix
prefix = Client().get_prefix(include_version=False)
t._registry[URL(prefix + 'o/', method='HEAD')] = Resp(
''.encode('utf-8'), 200, {}
)
t.register('/users/bob/personal_tokens/', json.dumps({
'token': 'abc123'
}), status_code=201, method='POST')
result = self.runner.invoke(
login, ['bob', '--password', 'secret', '--scope', 'read']
)
# Ensure that we got a zero exit status
self.assertEqual(result.exit_code, 0)
assert json.loads(t.requests[-1].body)['scope'] == 'read'
# Ensure that the output seems to be correct.
self.assertIn(mock.call(os.path.expanduser('~/.tower_cli.cfg'), 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('oauth_token = abc123\n'),
mock_open.mock_calls)
def test_sensor(self, mock_describe_job, hook_init, mock_client):
hook_init.return_value = None
mock_describe_job.side_effect = [
DESCRIBE_TRANSFORM_INPROGRESS_RESPONSE,
DESCRIBE_TRANSFORM_STOPPING_RESPONSE,
DESCRIBE_TRANSFORM_COMPELETED_RESPONSE
]
sensor = SageMakerTransformSensor(task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name')
sensor.execute(None)
# make sure we called 3 times(terminated when its compeleted)
self.assertEqual(mock_describe_job.call_count, 3)
# make sure the hook was initialized with the specific params
calls = [
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test')
]
hook_init.assert_has_calls(calls)
def test_sensor_with_log(self, mock_describe_job, mock_describe_job_with_log,
hook_init, mock_log_client, mock_client):
hook_init.return_value = None
mock_describe_job.return_value = DESCRIBE_TRAINING_COMPELETED_RESPONSE
mock_describe_job_with_log.side_effect = [
(LogState.WAIT_IN_PROGRESS, DESCRIBE_TRAINING_INPROGRESS_RESPONSE, 0),
(LogState.JOB_COMPLETE, DESCRIBE_TRAINING_STOPPING_RESPONSE, 0),
(LogState.COMPLETE, DESCRIBE_TRAINING_COMPELETED_RESPONSE, 0)
]
sensor = SageMakerTrainingSensor(
task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
print_log=True
)
sensor.execute(None)
self.assertEqual(mock_describe_job_with_log.call_count, 3)
self.assertEqual(mock_describe_job.call_count, 1)
calls = [
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test')
]
hook_init.assert_has_calls(calls)
def test_should_copy_files_to_subdirectory(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy):
# mock_get_conn.return_value =
source_bucket = self._create_bucket(name="SOURCE_BUCKET")
source_bucket.list_blobs.return_value = [
self._create_blob("FILE_A", "C1"),
self._create_blob("FILE_B", "C1"),
]
destination_bucket = self._create_bucket(name="DEST_BUCKET")
destination_bucket.list_blobs.return_value = []
mock_get_conn.return_value.bucket.side_effect = [source_bucket, destination_bucket]
self.gcs_hook.sync(
source_bucket="SOURCE_BUCKET", destination_bucket="DEST_BUCKET", destination_object="DEST_OBJ/"
)
mock_delete.assert_not_called()
mock_rewrite.assert_not_called()
mock_copy.assert_has_calls(
[
mock.call(
source_bucket="SOURCE_BUCKET",
source_object="FILE_A",
destination_bucket="DEST_BUCKET",
destination_object="DEST_OBJ/FILE_A",
),
mock.call(
source_bucket="SOURCE_BUCKET",
source_object="FILE_B",
destination_bucket="DEST_BUCKET",
destination_object="DEST_OBJ/FILE_B",
),
],
any_order=True,
)
def test_exec_success_csv_ensure_utc(self, gcs_hook_mock_class,
mysql_hook_mock_class):
"""Test successful run of execute function for CSV"""
op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID,
mysql_conn_id=MYSQL_CONN_ID,
sql=SQL,
export_format='CSV',
bucket=BUCKET,
filename=CSV_FILENAME,
ensure_utc=True)
mysql_hook_mock = mysql_hook_mock_class.return_value
mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS)
mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION
gcs_hook_mock = gcs_hook_mock_class.return_value
def _assert_upload(bucket, obj, tmp_filename, mime_type=None):
self.assertEqual(BUCKET, bucket)
self.assertEqual(CSV_FILENAME.format(0), obj)
self.assertEqual('text/csv', mime_type)
with open(tmp_filename, 'rb') as file:
self.assertEqual(b''.join(CSV_LINES), file.read())
gcs_hook_mock.upload.side_effect = _assert_upload
op.execute(None)
mysql_hook_mock_class.assert_called_once_with(
mysql_conn_id=MYSQL_CONN_ID)
mysql_hook_mock.get_conn().cursor().execute.assert_has_calls(
[mock.call(TZ_QUERY), mock.call(SQL)])
def test_wait_for_transfer_job(self, get_conn, mock_sleep):
list_method = get_conn.return_value.transferOperations.return_value.list
list_execute_method = list_method.return_value.execute
list_execute_method.side_effect = [
{OPERATIONS: [{METADATA: {STATUS: GcpTransferOperationStatus.IN_PROGRESS}}]},
{OPERATIONS: [{METADATA: {STATUS: GcpTransferOperationStatus.SUCCESS}}]},
]
get_conn.return_value.transferOperations.return_value.list_next.return_value = None
job_name = 'transferJobs/test-job'
self.gct_hook.wait_for_transfer_job({PROJECT_ID: TEST_PROJECT_ID, 'name': job_name})
self.assertTrue(list_method.called)
calls = [
mock.call(
filter=json.dumps({"project_id": TEST_PROJECT_ID, "job_names": [job_name]}),
name='transferOperations'
),
mock.call().execute(num_retries=5),
mock.call(
filter=json.dumps({"project_id": TEST_PROJECT_ID, "job_names": [job_name]}),
name='transferOperations'
),
mock.call().execute(num_retries=5)
]
list_method.assert_has_calls(calls)
args, kwargs = list_method.call_args_list[0]
self.assertEqual(
json.loads(kwargs['filter']),
{FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: ["transferJobs/test-job"]},
)
mock_sleep.assert_called_once_with(TIME_TO_SLEEP_IN_SECONDS)
def test_should_delete_extra_files(self, mock_get_conn, mock_delete,
mock_rewrite, mock_copy):
# mock_get_conn.return_value =
source_bucket = self._create_bucket(name="SOURCE_BUCKET")
source_bucket.list_blobs.return_value = []
destination_bucket = self._create_bucket(name="DEST_BUCKET")
destination_bucket.list_blobs.return_value = [
self._create_blob("SRC_OBJ/FILE_A", "C1", destination_bucket),
self._create_blob("SRC_OBJ/FILE_B", "C1", destination_bucket),
]
mock_get_conn.return_value.bucket.side_effect = [
source_bucket, destination_bucket
]
self.gcs_hook.sync(source_bucket="SOURCE_BUCKET",
destination_bucket="DEST_BUCKET",
delete_extra_files=True)
mock_delete.assert_has_calls(
[
mock.call("DEST_BUCKET", "SRC_OBJ/FILE_B"),
mock.call("DEST_BUCKET", "SRC_OBJ/FILE_A")
],
any_order=True,
)
mock_rewrite.assert_not_called()
mock_copy.assert_not_called()
def test_write_global_setting_deprecated(self):
"""Establish that if we attempt to write a valid setting, that
the parser's write method is run.
"""
# Invoke the command, but trap the file-write at the end
# so we don't plow over real things.
mock_open = mock.mock_open()
warning_text = 'The `--global` option is deprecated and will be '\
'removed. Use `--scope=global` to get the same effect.'
with mock.patch('tower_cli.cli.misc.open', mock_open,
create=True):
with mock.patch.object(os.path, 'isdir') as isdir:
with mock.patch.object(os, 'chmod'):
with mock.patch.object(warnings, 'warn') as warn:
isdir.return_value = True
result = self.runner.invoke(
config, ['username', 'meagan', '--global'],
)
warn.assert_called_once_with(warning_text,
DeprecationWarning)
self.assertEqual(warn.mock_calls[0][1][1],
DeprecationWarning)
isdir.assert_called_once_with('/etc/tower/')
# Ensure that the command completed successfully.
self.assertEqual(result.exit_code, 0)
self.assertEqual('Configuration updated successfully.',
result.output.strip())
# Ensure that the output seems to be correct.
self.assertIn(mock.call('/etc/tower/tower_cli.cfg', 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('username = meagan\n'),
mock_open.mock_calls)
def test_echo_setting(self):
"""Establish that the `echo_setting` method works in the way
that we expect.
"""
with settings.runtime_values(host='20.12.4.21'):
with mock.patch.object(click, 'secho') as secho:
_echo_setting('host')
self.assertEqual(secho.mock_calls, [
mock.call('host: ', fg='magenta', bold=True, nl=False),
mock.call('20.12.4.21', fg='white', bold=True),
])
def test_sync(self, run_task_mock):
run_task_mock.return_value = True
executor = DebugExecutor()
ti1 = MagicMock(key="t1")
ti2 = MagicMock(key="t2")
executor.tasks_to_run = [ti1, ti2]
executor.sync()
assert not executor.tasks_to_run
run_task_mock.assert_has_calls([mock.call(ti1), mock.call(ti2)])
def test_delete_objects(self, mock_hook):
operator = GoogleCloudStorageDeleteOperator(task_id=TASK_ID,
bucket_name=TEST_BUCKET,
objects=MOCK_FILES[0:2])
operator.execute(None)
mock_hook.return_value.list.assert_not_called()
mock_hook.return_value.delete.assert_has_calls(calls=[
mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[0]),
mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[1])
],
any_order=True)
def test_ensure_folders_exists_when_some_folders_exists(
self, mock_get_conn):
mock_get_conn.return_value.files.return_value.list.return_value.execute.side_effect = [
{
"files": [{
"id": "ID_1"
}]
},
{
"files": [{
"id": "ID_2"
}]
},
{
"files": []
},
]
mock_get_conn.return_value.files.return_value.create.return_value.execute.side_effect = [
{
"id": "ID_3"
},
{
"id": "ID_4"
},
]
result_value = self.gdrive_hook._ensure_folders_exists(
"AAA/BBB/CCC/DDD")
mock_get_conn.assert_has_calls(
[
mock.call().files().create(
body={
"name": "CCC",
"mimeType": "application/vnd.google-apps.folder",
"parents": ["ID_2"],
},
fields="id",
),
mock.call().files().create(
body={
"name": "DDD",
"mimeType": "application/vnd.google-apps.folder",
"parents": ["ID_3"],
},
fields="id",
),
],
any_order=True,
)
self.assertEqual("ID_4", result_value)
def test_read_from_file(self):
"""Give it some with '@' and test that it reads from the file"""
mock_open = mock.mock_open()
with mock.patch('tower_cli.utils.parser.open', mock_open, create=True):
manager = mock_open.return_value.__enter__.return_value
manager.read.return_value = 'foo: bar'
parser.process_extra_vars(["@fake_file1.yml"])
parser.process_extra_vars(["@fake_file2.yml", "@fake_file3.yml"])
# Ensure that "open" was triggered in test
self.assertIn(mock.call("fake_file1.yml", 'r'), mock_open.mock_calls)
self.assertIn(mock.call("fake_file2.yml", 'r'), mock_open.mock_calls)
self.assertIn(mock.call("fake_file3.yml", 'r'), mock_open.mock_calls)
def test_delete_prefix(self, mock_hook):
mock_hook.return_value.list.return_value = MOCK_FILES[1:3]
operator = GoogleCloudStorageDeleteOperator(task_id=TASK_ID,
bucket_name=TEST_BUCKET,
prefix=PREFIX)
operator.execute(None)
mock_hook.return_value.list.assert_called_once_with(
bucket_name=TEST_BUCKET, prefix=PREFIX)
mock_hook.return_value.delete.assert_has_calls(calls=[
mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[1]),
mock.call(bucket_name=TEST_BUCKET, object_name=MOCK_FILES[2])
],
any_order=True)
def test_read_from_file(self):
"""Give it some with '@' and test that it reads from the file"""
mock_open = mock.mock_open()
with mock.patch('tower_cli.utils.parser.open', mock_open, create=True):
manager = mock_open.return_value.__enter__.return_value
manager.read.return_value = 'foo: bar'
parser.extra_vars_loader_wrapper(["@fake_file1.yml"])
parser.extra_vars_loader_wrapper(["@fake_file2.yml",
"@fake_file3.yml"])
# Ensure that "open" was triggered in test
self.assertIn(mock.call("fake_file1.yml", 'r'), mock_open.mock_calls)
self.assertIn(mock.call("fake_file2.yml", 'r'), mock_open.mock_calls)
self.assertIn(mock.call("fake_file3.yml", 'r'), mock_open.mock_calls)
def test_execute(self, gcs_mock_hook, adls_one_mock_hook,
adls_two_mock_hook):
"""Test the execute function when the run is successful."""
operator = AdlsToGoogleCloudStorageOperator(
task_id=TASK_ID,
src_adls=ADLS_PATH_1,
dest_gcs=GCS_PATH,
replace=False,
azure_data_lake_conn_id=AZURE_CONN_ID,
google_cloud_storage_conn_id=GCS_CONN_ID)
adls_one_mock_hook.return_value.list.return_value = MOCK_FILES
adls_two_mock_hook.return_value.list.return_value = MOCK_FILES
# gcs_mock_hook.return_value.upload.side_effect = _assert_upload
uploaded_files = operator.execute(None)
gcs_mock_hook.return_value.upload.assert_has_calls([
mock.call(bucket_name='test',
filename=mock.ANY,
object_name='test/path/PARQUET.parquet',
gzip=False),
mock.call(bucket_name='test',
filename=mock.ANY,
object_name='test/path/TEST3.csv',
gzip=False),
mock.call(bucket_name='test',
filename=mock.ANY,
object_name='test/path/PIC.png',
gzip=False),
mock.call(bucket_name='test',
filename=mock.ANY,
object_name='test/TEST1.csv',
gzip=False),
mock.call(bucket_name='test',
filename=mock.ANY,
object_name='test/TEST2.csv',
gzip=False)
],
any_order=True)
adls_one_mock_hook.assert_called_once_with(
azure_data_lake_conn_id=AZURE_CONN_ID)
adls_two_mock_hook.assert_called_once_with(
azure_data_lake_conn_id=AZURE_CONN_ID)
gcs_mock_hook.assert_called_once_with(
google_cloud_storage_conn_id=GCS_CONN_ID, delegate_to=None)
# we expect MOCK_FILES to be uploaded
self.assertEqual(sorted(MOCK_FILES), sorted(uploaded_files))
def test_execute_wildcard_without_destination_object(self, mock_hook):
mock_hook.return_value.list.return_value = SOURCE_FILES_LIST
operator = GoogleCloudStorageToGoogleCloudStorageOperator(
task_id=TASK_ID, source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=DESTINATION_BUCKET)
operator.execute(None)
mock_calls_none = [
mock.call(TEST_BUCKET, 'test_object/file1.txt',
DESTINATION_BUCKET, 'test_object/file1.txt'),
mock.call(TEST_BUCKET, 'test_object/file2.txt',
DESTINATION_BUCKET, 'test_object/file2.txt'),
]
mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none)
def test_execute_wildcard_without_destination_object(self, mock_hook):
mock_hook.return_value.list.return_value = SOURCE_FILES_LIST
operator = GoogleCloudStorageToGoogleCloudStorageOperator(
task_id=TASK_ID, source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=DESTINATION_BUCKET)
operator.execute(None)
mock_calls_none = [
mock.call(TEST_BUCKET, 'test_object/file1.txt',
DESTINATION_BUCKET, 'test_object/file1.txt'),
mock.call(TEST_BUCKET, 'test_object/file2.txt',
DESTINATION_BUCKET, 'test_object/file2.txt'),
]
mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none)
def test_upsert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document(
{'data1': 'somedata'},
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document_id=test_id)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' +
self.test_collection_name, {
'data1': 'somedata',
'id': test_id
})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
def test_insert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosInsertDocumentOperator(
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document={
'id': test_id,
'data': 'sometestdata'
},
azure_cosmos_conn_id='azure_cosmos_test_key_id',
task_id='azure_cosmos_sensor')
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' +
self.test_collection_name, {
'data': 'sometestdata',
'id': test_id
})
]
self.cosmos.execute(None)
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_delete_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_database(self.test_database_name)
expected_calls = [mock.call().DeleteDatabase('dbs/test_database_name')]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_wc_with_last_modified_time_with_all_true_cond(self, mock_hook):
mock_hook.return_value.list.return_value = SOURCE_FILES_LIST
mock_hook.return_value.is_updated_after.side_effect = [True, True, True]
operator = GoogleCloudStorageToGoogleCloudStorageOperator(
task_id=TASK_ID, source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=DESTINATION_BUCKET,
last_modified_time=MOD_TIME_1)
operator.execute(None)
mock_calls_none = [
mock.call(TEST_BUCKET, 'test_object/file1.txt',
DESTINATION_BUCKET, 'test_object/file1.txt'),
mock.call(TEST_BUCKET, 'test_object/file2.txt',
DESTINATION_BUCKET, 'test_object/file2.txt'),
]
mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none)
def test_wait_for_job(self, mock_get_job):
mock_get_job.side_effect = [
mock.MagicMock(status=mock.MagicMock(state=JobStatus.RUNNING)),
mock.MagicMock(status=mock.MagicMock(state=JobStatus.ERROR)),
]
with self.assertRaises(AirflowException):
self.hook.wait_for_job(
job_id=JOB_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT,
wait_time=0,
)
calls = [
mock.call(location=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT),
mock.call(location=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT),
]
mock_get_job.has_calls(calls)
def test_wc_with_last_modified_time_with_all_true_cond(self, mock_hook):
mock_hook.return_value.list.return_value = SOURCE_FILES_LIST
mock_hook.return_value.is_updated_after.side_effect = [True, True, True]
operator = GoogleCloudStorageToGoogleCloudStorageOperator(
task_id=TASK_ID, source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=DESTINATION_BUCKET,
last_modified_time=MOD_TIME_1)
operator.execute(None)
mock_calls_none = [
mock.call(TEST_BUCKET, 'test_object/file1.txt',
DESTINATION_BUCKET, 'test_object/file1.txt'),
mock.call(TEST_BUCKET, 'test_object/file2.txt',
DESTINATION_BUCKET, 'test_object/file2.txt'),
]
mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none)
def test_create_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name)
expected_calls = [mock.call().CreateContainer(
'dbs/test_database_default',
{'id': self.test_collection_name})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_execute_wildcard_with_destination_object_retained_prefix(self, mock_hook):
mock_hook.return_value.list.return_value = SOURCE_FILES_LIST
operator = GoogleCloudStorageToGoogleCloudStorageOperator(
task_id=TASK_ID, source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=DESTINATION_BUCKET,
destination_object='{}/{}'.format(DESTINATION_OBJECT_PREFIX,
SOURCE_OBJECT_WILDCARD_SUFFIX[:-1])
)
operator.execute(None)
mock_calls_retained = [
mock.call(TEST_BUCKET, 'test_object/file1.txt',
DESTINATION_BUCKET, 'foo/bar/test_object/file1.txt'),
mock.call(TEST_BUCKET, 'test_object/file2.txt',
DESTINATION_BUCKET, 'foo/bar/test_object/file2.txt'),
]
mock_hook.return_value.rewrite.assert_has_calls(mock_calls_retained)
def test_public_application_scoped_token(self):
"""Establish that if `tower-cli login` is called with a username,
password, and public OAuth2 client ID, we obtain and write an oauth
token to the config file
"""
# Invoke the command.
mock_open = mock.mock_open()
with mock.patch('tower_cli.cli.misc.open', mock_open,
create=True):
with mock.patch.object(os, 'chmod'):
with client.test_mode as t:
# You have to modify this internal private registry to
# register a URL endpoint that _doesn't_ have the version
# prefix
prefix = Client().get_prefix(include_version=False)
t._registry[URL(prefix + 'o/', method='HEAD')] = Resp(
''.encode('utf-8'), 200, {}
)
t._registry[URL(prefix + 'o/token/', method='POST')] = Resp(
json.dumps({'access_token': 'abc123'}).encode('utf-8'),
201, {}
)
result = self.runner.invoke(
login, ['bob', '--password', 'secret', '--client-id',
'abc123']
)
# Ensure that we got a zero exit status
self.assertEqual(result.exit_code, 0)
data = urlparse.parse_qs(t.requests[-1].body)
assert data['scope'] == ['write']
assert data['grant_type'] == ['password']
assert data['password'] == ['secret']
assert data['username'] == ['bob']
assert data['client_id'] == ['abc123']
# Ensure that the output seems to be correct.
self.assertIn(mock.call(os.path.expanduser('~/.tower_cli.cfg'), 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('oauth_token = abc123\n'),
mock_open.mock_calls)
def test_upsert_document_default(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document({'id': test_id})
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'id': test_id})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
def test_write_setting(self):
"""Establish that if we attempt to write a valid setting, that
the parser's write method is run.
"""
# Invoke the command, but trap the file-write at the end
# so we don't plow over real things.
mock_open = mock.mock_open()
with mock.patch('tower_cli.commands.config.open', mock_open,
create=True):
result = self.runner.invoke(config, ['username', 'luke'])
# Ensure that the command completed successfully.
self.assertEqual(result.exit_code, 0)
self.assertEqual(result.output.strip(),
'Configuration updated successfully.')
# Ensure that the output seems to be correct.
self.assertIn(mock.call(os.path.expanduser('~/.tower_cli.cfg'), 'w'),
mock_open.mock_calls)
self.assertIn(mock.call().write('username = luke\n'),
mock_open.mock_calls)
def test_insert_documents(self, cosmos_mock):
test_id1 = str(uuid.uuid4())
test_id2 = str(uuid.uuid4())
test_id3 = str(uuid.uuid4())
documents = [
{'id': test_id1, 'data': 'data1'},
{'id': test_id2, 'data': 'data2'},
{'id': test_id3, 'data': 'data3'}]
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.insert_documents(documents)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data1', 'id': test_id1}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data2', 'id': test_id2}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data3', 'id': test_id3})]
logging.getLogger().info(returned_item)
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_insert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosInsertDocumentOperator(
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document={'id': test_id, 'data': 'sometestdata'},
azure_cosmos_conn_id='azure_cosmos_test_key_id',
task_id='azure_cosmos_sensor')
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' + self.test_collection_name,
{'data': 'sometestdata', 'id': test_id})]
self.cosmos.execute(None)
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_unset(self):
"""Establish that calling `tower-cli config --unset` works in the
way that we expect.
"""
# Invoke the command, but trap the file-write at the end
# so we don't plow over real things.
mock_open = mock.mock_open()
with mock.patch('tower_cli.commands.config.open', mock_open,
create=True):
result = self.runner.invoke(config, ['username', '--unset'])
# Ensure that the command completed successfully.
self.assertEqual(result.exit_code, 0)
self.assertEqual(result.output.strip(),
'Configuration updated successfully.')
# Ensure that the output seems to be correct.
self.assertNotIn(mock.call().write('username = luke\n'),
mock_open.mock_calls)
def test_delete_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_database(self.test_database_name)
expected_calls = [mock.call().DeleteDatabase('dbs/test_database_name')]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
