def test_session_creation_sensor(mocker):
sen = LivySessionCreationSensor(
session_id=123,
poke_interval=1,
timeout=3,
task_id="test_session_creation_sensor",
)
response = mock_http_calls(
200,
content=b'{"id": 123, "state": "idle"}',
)
mocker.patch.object(HttpHook, "get_conn", return_value=response)
sen.execute({})
assert sen.session_id == 123
def test_session_creation_sensor_unknown_state(mocker):
sen = LivySessionCreationSensor(
session_id=123,
poke_interval=2,
timeout=5,
task_id="test_session_creation_sensor_malformed_json",
)
http_response = mock_http_calls(200,
content=b'{"id": 123, "state": "IDK"}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowException) as ae:
sen.execute({})
print(f"Imitated unknown session state for session creation sensor, "
f"got the expected exception:\n<{ae.value}>\n")
def test_session_creation_sensor_timeout(mocker):
sen = LivySessionCreationSensor(
session_id=123,
poke_interval=1,
timeout=3,
task_id="test_session_creation_sensor_timeout",
)
http_response = mock_http_calls(
200, content=b'{"id": 123, "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 session "
f"to be created, got the expected exception:\n<{te.value}>\n"
f"API was polled {http_response.send.call_count} times.")
def test_session_creation_sensor_invalid_timings(poke_interval, timeout):
print(f"\n\nRunning test_session_creation_sensor_invalid_timings with:\n"
f"poke_interval={poke_interval}, timeout={timeout}")
with raises(AirflowException) as ae:
LivySessionCreationSensor(
session_id=123,
poke_interval=poke_interval,
timeout=timeout,
task_id="test_session_creation_sensor_invalid_timings",
)
print(f"Set up the session creation sensor with invalid timings, "
f"got the expected exception:\n<{ae.value}>\n")