def test_wait_for_completion(_mock_is_emr_step_complete, _mock_read_events):
launcher = EmrPySparkStepLauncher(
region_name="",
staging_bucket="",
staging_prefix="",
wait_for_logs=False,
action_on_failure="",
cluster_id="",
spark_config={},
local_pipeline_package_path="",
deploy_local_pipeline_package=False,
)
yielded_events = list(
launcher.wait_for_completion(mock.MagicMock(), None, None, None, None, check_interval=0)
)
assert yielded_events == EVENTS
def test_fetch_logs_on_fail(_mock_log_step_event, mock_log_logs,
mock_wait_for_completion, _mock_boto3_resource):
mock_log = mock.MagicMock()
mock_wait_for_completion.side_effect = EmrError()
step_launcher = EmrPySparkStepLauncher(
region_name="",
staging_bucket="",
staging_prefix="",
spark_config=None,
action_on_failure="",
cluster_id="",
local_pipeline_package_path="",
deploy_local_pipeline_package=False,
wait_for_logs=True,
)
with pytest.raises(EmrError):
for _ in step_launcher.wait_for_completion_and_log(
mock_log, None, None, None, None):
pass
assert mock_log_logs.call_count == 1