def test_required_args():
waiter = AWSClientWait(waiter_name="JobRunning")
with pytest.raises(ValueError, match="An AWS client string"):
waiter.run()
waiter = AWSClientWait(client="batch")
with pytest.raises(ValueError, match="A waiter name"):
waiter.run()
def test_prefect_batch_waiters_success(monkeypatch, boto_client):
batch_wait_task = AWSClientWait(client="batch")
create_waiter_mock = MagicMock()
monkeypatch.setattr(
"prefect.tasks.aws.client_waiter.create_waiter_with_client",
create_waiter_mock)
for i, waiter_name in enumerate(["JobExists", "JobComplete",
"JobRunning"]):
batch_wait_task.run(waiter_name=waiter_name)
args, _ = create_waiter_mock.call_args_list[i]
assert args[0] == waiter_name
assert isinstance(args[1], WaiterModel)
assert len(args[1].waiter_names) == 3
assert args[1].version == 2
def test_boto_waiter(boto_client):
boto_client.waiter_names = ["BotoWaiter"]
waiter = MagicMock()
boto_client.get_waiter = MagicMock(return_value=waiter)
batch_wait_task = AWSClientWait(client="batch")
batch_wait_task.run(
waiter_name="BotoWaiter",
waiter_kwargs={"WaiterConfig": {
"Delay": 20,
"MaxAttempts": 300
}},
)
# Check waiter fetched from boto client
args, _ = boto_client.get_waiter.call_args_list[0]
assert args[0] == "BotoWaiter"
# Check
_, kwargs = waiter.wait.call_args_list[0]
assert kwargs["WaiterConfig"] == {"Delay": 20, "MaxAttempts": 300}
def test_waiter_custom_definition(monkeypatch, boto_client):
waiter = MagicMock()
monkeypatch.setattr(
"prefect.tasks.aws.client_waiter.create_waiter_with_client",
MagicMock(return_value=waiter),
)
batch_wait_task = AWSClientWait(client="batch",
waiter_definition=TEST_WAITERS)
batch_wait_task.run(waiter_name="TestWaiter")
assert waiter.wait.called
batch_wait_task.run(
waiter_name="TestWaiter",
waiter_kwargs={"WaiterConfig": {
"Delay": 20,
"MaxAttempts": 300
}},
)
_, kwargs = waiter.wait.call_args_list[1]
assert kwargs["WaiterConfig"] == {"Delay": 20, "MaxAttempts": 300}
def test_prefect_waiter_fail(boto_client):
batch_wait_task = AWSClientWait(client="batch")
with pytest.raises(ValueError,
match="Unable to load waiter 'UnknownWaiter'"):
batch_wait_task.run(waiter_name="UnknownWaiter")