def test_batch_sensor(mocker):
sen = LivyBatchSensor(batch_id=2, task_id="test_batch_sensor")
http_response = mock_http_calls(
200,
content=b'{"id": 2, "state": "success"}',
)
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
sen.execute({})
def test_batch_sensor_timeout(mocker):
sen = LivyBatchSensor(
batch_id=2,
task_id="test_batch_sensor_timeout",
poke_interval=1,
timeout=2,
)
http_response = mock_http_calls(200,
content=b'{"id": 2, "state": "starting"}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowSensorTimeout) as te:
sen.execute({})
assert 2 <= http_response.send.call_count <= 4
print(
f"\n\nImitated period=1sec, timeout=2sec while waiting for batch to finish, "
f"got the expected exception:\n<{te.value}>\n"
f"API was polled {http_response.send.call_count} times.")