def test_it_defaults_wait_duration(sqs_mock, read_queue_mock, sf_client_mock):
sqs_mock.Queue.return_value = sqs_mock
sf_client_mock.start_execution.return_value = execution_stub()
read_queue_mock.return_value = [
SimpleNamespace(
body=json.dumps({
"hello": "world",
"QueryExecutor": "athena"
}),
receipt_handle="1234",
)
]
expected_call = json.dumps({
"hello": "world",
"QueryExecutor": "athena",
"AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID": "1234",
"JobId": "4231",
"WaitDuration": 15
})
handler({
"ExecutionId": "1234",
"ExecutionName": "4231",
}, SimpleNamespace())
sf_client_mock.start_execution.assert_called_with(stateMachineArn=ANY,
input=expected_call)
def test_it_recognises_completed_executions(clear_mock, load_mock, sqs_mock,
read_queue_mock, sf_client_mock):
sqs_mock.Queue.return_value = sqs_mock
sf_client_mock.start_execution.return_value = execution_stub()
read_queue_mock.return_value = []
load_mock.side_effect = [
*[execution_stub(status="SUCCEEDED") for _ in range(0, 10)],
*[
execution_stub(status="RUNNING", ReceiptHandle="handle")
for _ in range(10, 15)
],
]
handler(
{
"ExecutionId": "1234",
"ExecutionName": "4231",
"RunningExecutions": {
"Data": list(range(0, 15)),
"Total": 20
},
},
SimpleNamespace(),
)
read_queue_mock.assert_called_with(ANY, 10)
def test_it_abandons_when_any_query_fails_and_no_running_in_current_loop(
mock_abandon, mock_load):
mock_load.return_value = execution_stub(status="FAILED")
mock_abandon.side_effect = RuntimeError
with pytest.raises(RuntimeError):
handler(
{
"ExecutionId": "1234",
"ExecutionName": "4321",
"RunningExecutions": {
"IsFailing": True,
"Data": [{}],
"Total": 1,
}
}, SimpleNamespace())
def test_it_starts_machine_as_expected(sqs_mock, read_queue_mock,
sf_client_mock):
sqs_mock.Queue.return_value = sqs_mock
sf_client_mock.start_execution.return_value = execution_stub()
read_queue_mock.return_value = [
SimpleNamespace(
body=json.dumps({
"hello": "world",
"QueryExecutor": "athena"
}),
receipt_handle="1234",
)
]
expected_call = json.dumps({
"hello": "world",
"QueryExecutor": "athena",
"AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID": "1234",
"JobId": "4231",
"WaitDuration": 5
})
resp = handler(
{
"ExecutionId": "1234",
"ExecutionName": "4231",
"AthenaConcurrencyLimit": 5,
"QueryExecutionWaitSeconds": 5,
}, SimpleNamespace())
sf_client_mock.start_execution.assert_called_with(stateMachineArn=ANY,
input=expected_call)
assert 1 == resp["Total"]
assert not resp["IsFailing"]
def test_it_abandons_when_previous_loop_found_failure(mock_clear, mock_abandon,
mock_load):
mock_load.return_value = execution_stub(status="SUCCEEDED",
ReceiptHandle="handle")
mock_abandon.side_effect = RuntimeError
with pytest.raises(RuntimeError):
handler(
{
"ExecutionId": "1234",
"ExecutionName": "4321",
"RunningExecutions": {
"IsFailing": True,
"Data": [{}],
"Total": 1,
}
}, SimpleNamespace())
def test_it_starts_state_machine_per_message(sqs_mock, read_queue_mock,
sf_client_mock):
sqs_mock.Queue.return_value = sqs_mock
sf_client_mock.start_execution.return_value = execution_stub()
read_queue_mock.return_value = [
SimpleNamespace(
body=json.dumps({
"hello": "world",
"QueryExecutor": "athena"
}),
receipt_handle="1234",
),
SimpleNamespace(
body=json.dumps({
"other": "world",
"QueryExecutor": "athena"
}),
receipt_handle="4321",
)
]
resp = handler({
"ExecutionId": "1234",
"ExecutionName": "4231",
}, SimpleNamespace())
assert 2 == sf_client_mock.start_execution.call_count
assert 2 == resp["Total"]
assert not resp["IsFailing"]
def test_raises_error_for_invalid_executor(sqs_mock, read_queue_mock,
sf_client_mock):
sqs_mock.Queue.return_value = sqs_mock
sf_client_mock.start_execution.return_value = execution_stub()
read_queue_mock.return_value = [
SimpleNamespace(
body=json.dumps({
"hello": "world",
"QueryExecutor": "unknown"
}),
receipt_handle="1234",
)
]
with pytest.raises(NotImplementedError):
handler({
"ExecutionId": "1234",
"ExecutionName": "4231",
}, SimpleNamespace())
def test_limits_calls_to_capacity(sqs_mock, read_queue_mock, sf_client_mock):
sqs_mock.Queue.return_value = sqs_mock
sf_client_mock.start_execution.return_value = execution_stub()
read_queue_mock.return_value = [
SimpleNamespace(
body=json.dumps({
"hello": "world",
"QueryExecutor": "athena"
}),
receipt_handle="1234",
),
]
handler(
{
"ExecutionId": "1234",
"ExecutionName": "4231",
"AthenaConcurrencyLimit": 20,
}, SimpleNamespace())
read_queue_mock.assert_called_with(ANY, 20)
def test_it_skips_with_no_remaining_capacity(mock_load, sqs_mock,
read_queue_mock):
sqs_mock.Queue.return_value = sqs_mock
mock_load.return_value = execution_stub(status="RUNNING",
ReceiptHandle="handle")
resp = handler(
{
"ExecutionId": "1234",
"ExecutionName": "4231",
"RunningExecutions": {
"Data": list(range(0, 20)),
"Total": 20
}
}, SimpleNamespace())
read_queue_mock.assert_not_called()
assert 20 == resp["Total"]
assert not resp["IsFailing"]
def test_it_waits_for_running_executions_before_abandoning(
mock_sf, mock_abandon, mock_load):
mock_load.side_effect = [
execution_stub(status="FAILED", ReceiptHandle="handle1"),
execution_stub(status="RUNNING", ReceiptHandle="handle2")
]
res = handler(
{
"ExecutionId": "1234",
"ExecutionName": "4321",
"RunningExecutions": {
"IsFailing": False,
"Data": [{}, {}],
"Total": 2,
}
}, SimpleNamespace())
mock_abandon.assert_not_called()
mock_sf.start_execution.assert_not_called()
assert 1 == res["Total"]
assert res["IsFailing"]
